Upload folder using huggingface_hub

modeling/__init__.py

@@ -0,0 +1,5 @@
+from .embedding_model import DeepXEmbeddingModel
+from .gdn2_attention import GatedDeltaNet2Attention
+from .pipeline import DeepXPipeline
+
+__all__ = ["DeepXEmbeddingModel", "GatedDeltaNet2Attention", "DeepXPipeline"]

modeling/embedding_model.py

@@ -0,0 +1,200 @@
+"""
+DeepX v0.7: Gated DeltaNet-2 Hyperloop Backbone + ColBERT Head.
+
+Architecture:
+  Begin(4 NarrowA) → Phase1×2 [WideA + NarrowA×4] → Phase2×4 [NarrowB×4 + WideB]
+  → End(1 WideB) = 35 compute passes
+
+Unique layers: 9 (4 begin + 4 shared cores + 1 end)
+Per-loop differentiation: LoRA + RoDE
+
+Outputs:
+  1. Single vector (1536-d) via attention pooling
+  2. Token vectors (T × 128-d) via ColBERT head
+
+Weight Init: ~95% from Gemma 4 E2B via direct copy + SVD LoRA.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import logging
+from typing import Optional, Tuple, Union
+
+from config import DeepXConfig
+from .hybrid_layer import (
+    DeepXLayer, make_narrow_a_layer, make_narrow_b_layer,
+    make_wide_a_layer, make_wide_b_layer,
+)
+from .hyperloop import HyperloopPhase
+from .utils import RMSNorm, RoDE
+
+logger = logging.getLogger(__name__)
+
+
+class ColBERTHead(nn.Module):
+    """Projects hidden states to low-dimensional token vectors for MaxSim."""
+
+    def __init__(self, hidden_size: int, colbert_dim: int = 128):
+        super().__init__()
+        self.linear = nn.Linear(hidden_size, colbert_dim, bias=False)
+
+    def forward(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        token_embeds = self.linear(hidden_states)
+        token_embeds = F.normalize(token_embeds, p=2, dim=-1)
+        if attention_mask is not None:
+            token_embeds = token_embeds * attention_mask.unsqueeze(-1).to(token_embeds.dtype)
+        return token_embeds
+
+
+class DeepXEmbeddingModel(nn.Module):
+    """
+    DeepX v0.7 Backbone — Gated DeltaNet-2 Hyperloop + ColBERT.
+    
+    Receives hidden_states from external frozen token embedding.
+    """
+
+    def __init__(self, config: DeepXConfig):
+        super().__init__()
+        self.config = config
+
+        # ═══ 1. Begin Block: 4 unique NarrowA layers (direct copy from Gemma 0-3) ═══
+        self.begin_blocks = nn.ModuleList([
+            make_narrow_a_layer(config, layer_idx=i)
+            for i in range(config.begin_layers)
+        ])
+
+        # ═══ 2. Phase1 Loop: 2 iterations × [WideA + NarrowA×4] ═══
+        self.shared_narrow_a = make_narrow_a_layer(config, layer_idx=config.begin_layers)
+        self.shared_wide_a = make_wide_a_layer(config, layer_idx=config.begin_layers + 1)
+
+        # Attach RoDE to shared cores
+        if config.use_rode:
+            self.shared_narrow_a.self_attn._rode = RoDE(
+                dim=config.depth_rotary_dims, num_loops=config.phase1_loops
+            )
+            self.shared_wide_a.self_attn._rode = RoDE(
+                dim=config.depth_rotary_dims, num_loops=config.phase1_loops
+            )
+
+        self.phase1 = HyperloopPhase(
+            config=config,
+            shared_narrow=self.shared_narrow_a,
+            shared_wide=self.shared_wide_a,
+            num_loops=config.phase1_loops,
+            narrow_num_heads=config.narrow_a_heads,
+            narrow_kv_heads=config.narrow_a_kv_heads,
+            narrow_head_dim=config.narrow_a_head_dim,
+            narrow_intermediate=config.narrow_a_intermediate,
+            wide_num_heads=config.wide_a_heads,
+            wide_kv_heads=config.wide_a_kv_heads,
+            wide_head_dim=config.wide_a_head_dim,
+            wide_intermediate=config.wide_a_intermediate,
+            wide_first=True,  # [WideA, NarrowA×4]
+        )
+
+        # ═══ 3. Phase2 Loop: 4 iterations × [NarrowB×4 + WideB] ═══
+        self.shared_narrow_b = make_narrow_b_layer(config, layer_idx=config.begin_layers + 2)
+        self.shared_wide_b = make_wide_b_layer(config, layer_idx=config.begin_layers + 3)
+
+        if config.use_rode:
+            self.shared_narrow_b.self_attn._rode = RoDE(
+                dim=config.depth_rotary_dims, num_loops=config.phase2_loops
+            )
+            self.shared_wide_b.self_attn._rode = RoDE(
+                dim=config.depth_rotary_dims, num_loops=config.phase2_loops
+            )
+
+        self.phase2 = HyperloopPhase(
+            config=config,
+            shared_narrow=self.shared_narrow_b,
+            shared_wide=self.shared_wide_b,
+            num_loops=config.phase2_loops,
+            narrow_num_heads=config.narrow_b_heads,
+            narrow_kv_heads=config.narrow_b_kv_heads,
+            narrow_head_dim=config.narrow_b_head_dim,
+            narrow_intermediate=config.narrow_b_intermediate,
+            wide_num_heads=config.wide_b_heads,
+            wide_kv_heads=config.wide_b_kv_heads,
+            wide_head_dim=config.wide_b_head_dim,
+            wide_intermediate=config.wide_b_intermediate,
+            wide_first=False,  # [NarrowB×4, WideB]
+        )
+
+        # ═══ 4. End Block: 1 unique WideB layer (layer 34 direct copy) ═══
+        self.end_block = make_wide_b_layer(config, layer_idx=config.begin_layers + 4)
+
+        # ═══ Final norm ═══
+        self.norm = RMSNorm(config.hidden_size, config.rms_norm_eps)
+
+        # ═══ Output Head 1: Attention Pooling ═══
+        self.pooling_strategy = config.pooling_strategy
+        if config.pooling_strategy == "attention":
+            self.pool_query = nn.Parameter(torch.randn(1, 1, config.hidden_size) * 0.02)
+
+        # ═══ Output Head 2: ColBERT ═══
+        self.use_colbert = config.use_colbert
+        if config.use_colbert:
+            self.colbert_head = ColBERTHead(config.hidden_size, config.colbert_dim)
+
+        self.to(config.torch_dtype)
+
+    def _pool(self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor:
+        if self.pooling_strategy == "attention":
+            scale = hidden_states.shape[-1] ** -0.5
+            scores = (hidden_states * self.pool_query).sum(dim=-1) * scale
+            if attention_mask is not None:
+                scores = scores.masked_fill(attention_mask == 0, float("-inf"))
+            weights = F.softmax(scores, dim=-1).unsqueeze(-1)
+            return (hidden_states * weights).sum(dim=1)
+        # Mean pooling fallback
+        if attention_mask is not None:
+            mask = attention_mask.unsqueeze(-1).to(hidden_states.dtype)
+            return (hidden_states * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)
+        return hidden_states.mean(dim=1)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        normalize: bool = True,
+        truncate_dim: Optional[int] = None,
+        return_colbert: bool = False,
+    ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
+        B, T, _ = hidden_states.shape
+        position_ids = torch.arange(T, device=hidden_states.device).unsqueeze(0).expand(B, -1)
+
+        # 1. Begin blocks (4 unique NarrowA)
+        for layer in self.begin_blocks:
+            hidden_states = layer(hidden_states, attention_mask=attention_mask, position_ids=position_ids)
+
+        # 2. Phase1 loop: 2 × [WideA + NarrowA×4] = 10 passes
+        hidden_states = self.phase1(hidden_states, attention_mask=attention_mask, position_ids=position_ids)
+
+        # 3. Phase2 loop: 4 × [NarrowB×4 + WideB] = 20 passes
+        hidden_states = self.phase2(hidden_states, attention_mask=attention_mask, position_ids=position_ids)
+
+        # 4. End block (1 unique WideB)
+        hidden_states = self.end_block(hidden_states, attention_mask=attention_mask, position_ids=position_ids)
+
+        # 5. Final norm
+        hidden_states = self.norm(hidden_states)
+
+        # --- Output 1: Single vector ---
+        single_embed = self._pool(hidden_states, attention_mask)
+        if truncate_dim is not None:
+            single_embed = single_embed[:, :truncate_dim]
+        if normalize:
+            single_embed = F.normalize(single_embed, p=2, dim=-1)
+
+        # --- Output 2: ColBERT token vectors ---
+        if return_colbert and self.use_colbert:
+            token_embeds = self.colbert_head(hidden_states, attention_mask)
+            return single_embed, token_embeds
+
+        return single_embed
+
+    def count_parameters(self) -> dict:
+        total = sum(p.numel() for p in self.parameters())
+        trainable = sum(p.numel() for p in self.parameters() if p.requires_grad)
+        return {"backbone_total": total, "trainable": trainable}

modeling/gdn2_attention.py

@@ -0,0 +1,411 @@
+"""
+Gated DeltaNet-2 Dual-Path Attention.
+
+Dual-path design:
+  Path 1 (Softmax): Standard chunked softmax attention with RoPE
+          → Uses Gemma Q,K,V weights directly, 100% compatible
+  Path 2 (GDN-2):  Gated Delta Rule-2 with channel-wise erase/write gates
+          → Learns linear attention gradually (init as near-no-op)
+
+Merge: output = α × softmax_path + (1-α) × gdn2_path
+  α init ≈ 1.0 → model starts as pure softmax (Gemma weights work perfectly)
+  α learned → shifts to GDN-2 as training progresses (O(T) inference)
+
+Gated Delta Rule-2 (per timestep):
+  1. Decay:  A_t = diag(α_t) @ A_{t-1}
+  2. Erase:  u_t = b_t ⊙ k_t  (channel-wise erase gate)
+  3. Write:  A_t -= u_t^T @ (u_t @ A_{t-1} - w_t ⊙ v_t)
+  4. Read:   o_t = q_t @ A_t
+
+Reference: Hatamizadeh et al. "Gated DeltaNet-2" (NVIDIA, May 2026)
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from typing import Optional
+from einops import rearrange
+
+from .utils import RMSNorm, YaRNRotaryEmbedding, apply_rotary_pos_emb, apply_depth_rotary_emb
+
+
+class ShortConv1d(nn.Module):
+    """Causal 1D convolution for Q,K preprocessing (init as identity pass-through)."""
+    
+    def __init__(self, dim: int, kernel_size: int = 4):
+        super().__init__()
+        self.conv = nn.Conv1d(dim, dim, kernel_size, padding=kernel_size - 1, groups=dim)
+        # Init as identity: center weight = 1, rest = 0
+        nn.init.zeros_(self.conv.weight)
+        nn.init.zeros_(self.conv.bias)
+        # Set last position (causal) to 1 for identity
+        with torch.no_grad():
+            self.conv.weight[:, :, -1] = 1.0
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """x: (B, T, D) → (B, T, D)"""
+        x = x.transpose(1, 2)  # (B, D, T)
+        x = self.conv(x)[..., :x.shape[-1]]  # causal: trim future padding
+        return x.transpose(1, 2)  # (B, T, D)
+
+
+class GatedDeltaNet2Attention(nn.Module):
+    """
+    Dual-Path: Softmax + Gated DeltaNet-2.
+    
+    Init strategy: α ≈ 1 (pure softmax) → Gemma weights work immediately.
+    Training shifts α towards GDN-2 path for O(T) inference.
+    """
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        head_dim: int,
+        chunk_size: int = 64,
+        attention_dropout: float = 0.0,
+        layer_idx: int = 0,
+        use_dual_path: bool = True,
+        softmax_init_alpha: float = 5.0,
+        use_short_conv: bool = True,
+        conv_kernel_size: int = 4,
+        # RoPE params
+        max_position_embeddings: int = 131072,
+        rope_theta: float = 1000000.0,
+        rope_scaling_factor: float = 32.0,
+        rope_original_max_position: int = 4096,
+        rope_beta_fast: int = 32,
+        rope_beta_slow: int = 1,
+    ):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_dim = head_dim
+        self.num_kv_groups = num_heads // num_kv_heads
+        self.layer_idx = layer_idx
+        self.chunk_size = chunk_size
+        self.use_dual_path = use_dual_path
+
+        # === Shared projections (from Gemma 4 E2B) ===
+        self.q_proj = nn.Linear(hidden_size, num_heads * head_dim, bias=False)
+        self.k_proj = nn.Linear(hidden_size, num_kv_heads * head_dim, bias=False)
+        self.v_proj = nn.Linear(hidden_size, num_kv_heads * head_dim, bias=False)
+        self.o_proj = nn.Linear(num_heads * head_dim, hidden_size, bias=False)
+
+        # === Q/K normalization (L2 norm per head — GDN-2 best practice) ===
+        self.q_norm = RMSNorm(head_dim)
+        self.k_norm = RMSNorm(head_dim)
+
+        # === RoPE for softmax path ===
+        self.rotary_emb = YaRNRotaryEmbedding(
+            dim=head_dim,
+            max_position_embeddings=max_position_embeddings,
+            base=rope_theta,
+            scaling_factor=rope_scaling_factor,
+            original_max_position=rope_original_max_position,
+            beta_fast=rope_beta_fast,
+            beta_slow=rope_beta_slow,
+        )
+
+        # === Short conv for GDN-2 path Q,K (init as identity) ===
+        if use_short_conv:
+            self.q_conv = ShortConv1d(num_heads * head_dim, conv_kernel_size)
+            self.k_conv = ShortConv1d(num_kv_heads * head_dim, conv_kernel_size)
+        else:
+            self.q_conv = None
+            self.k_conv = None
+
+        # === GDN-2 specific params (all NEW, init for near-no-op) ===
+        # Erase gate: b_t = sigmoid(W_b @ x_t) ∈ [0,1]^{d_k}
+        self.erase_proj = nn.Linear(hidden_size, num_kv_heads * head_dim, bias=True)
+        nn.init.zeros_(self.erase_proj.weight)
+        nn.init.zeros_(self.erase_proj.bias)  # sigmoid(0)=0.5 → moderate erase
+        
+        # Write gate: w_t = sigmoid(W_w @ x_t) ∈ [0,1]^{d_v}  
+        self.write_proj = nn.Linear(hidden_size, num_kv_heads * head_dim, bias=True)
+        nn.init.zeros_(self.write_proj.weight)
+        nn.init.zeros_(self.write_proj.bias)  # sigmoid(0)=0.5 → moderate write
+        
+        # Decay (channel-wise, log-parameterized for stability)
+        # g_t = -exp(a) * softplus(W_f @ x_t + δ)
+        # α_t = exp(g_t) ∈ (0, 1]
+        # Init: a=-5, δ=5 → g≈0 → α≈1 (no decay initially)
+        self.decay_a = nn.Parameter(torch.full((num_kv_heads * head_dim,), -5.0))
+        self.decay_delta = nn.Parameter(torch.full((num_kv_heads * head_dim,), 5.0))
+        self.decay_proj = nn.Linear(hidden_size, num_kv_heads * head_dim, bias=False)
+        nn.init.zeros_(self.decay_proj.weight)  # → softplus(0+5)≈5, g=-exp(-5)*5≈-0.03, α≈0.97
+
+        # === Dual-path mix: α (softmax weight) ===
+        # Init large positive → sigmoid ≈ 1 → pure softmax at start
+        self.path_mix_logit = nn.Parameter(torch.full((num_heads,), softmax_init_alpha))
+
+        # === Output gate (SiLU gating like GDN-2 paper) ===
+        self.output_gate_proj = nn.Linear(hidden_size, num_heads * head_dim, bias=True)
+
+        # === RoDE placeholder ===
+        self._rode = None
+
+        self.dropout = nn.Dropout(attention_dropout)
+
+    def _expand_kv(self, x: torch.Tensor) -> torch.Tensor:
+        """Expand KV heads for GQA: (B, H_kv, T, D) → (B, H, T, D)"""
+        if self.num_kv_groups == 1:
+            return x
+        B, H_kv, T, D = x.shape
+        x = x.unsqueeze(2).expand(-1, -1, self.num_kv_groups, -1, -1)
+        return x.reshape(B, self.num_heads, T, D)
+
+    # ── Path 1: Standard Softmax Attention (uses RoPE) ──────────
+
+    def _softmax_attention(
+        self,
+        q: torch.Tensor,      # (B, H, T, D)
+        k: torch.Tensor,      # (B, H, T, D) — already expanded for GQA
+        v: torch.Tensor,      # (B, H, T, D)
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        B, H, T, D = q.shape
+        scale = 1.0 / math.sqrt(D)
+        
+        # Full attention for short seqs, chunked for long
+        scores = torch.matmul(q, k.transpose(-1, -2)) * scale  # (B, H, T, T)
+        
+        if attention_mask is not None:
+            # attention_mask: (B, T) → (B, 1, 1, T)
+            mask = attention_mask[:, None, None, :].to(scores.dtype)
+            scores = scores.masked_fill(mask == 0, float("-inf"))
+        
+        attn_weights = F.softmax(scores, dim=-1)
+        attn_weights = torch.nan_to_num(attn_weights, nan=0.0)
+        attn_weights = self.dropout(attn_weights)
+        
+        return torch.matmul(attn_weights, v)
+
+    # ── Path 2: Gated DeltaNet-2 (Recurrent, O(T) inference) ───
+
+    def _gdn2_attention(
+        self,
+        q: torch.Tensor,      # (B, H, T, D) — L2 normalized
+        k: torch.Tensor,      # (B, H_kv, T, D) — L2 normalized  
+        v: torch.Tensor,      # (B, H_kv, T, D)
+        erase_gate: torch.Tensor,  # (B, T, d_kv)
+        write_gate: torch.Tensor,  # (B, T, d_kv)
+        decay: torch.Tensor,       # (B, T, d_kv) — α_t values
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """
+        Gated Delta Rule-2 via FLA chunk-parallel kernel (O(T) training).
+        Falls back to sequential recurrence if FLA not available.
+        """
+        B, H_kv, T, D = k.shape
+        H = self.num_heads
+        dtype = q.dtype
+
+        # Expand KV for GQA
+        k_exp = self._expand_kv(k)        # (B, H, T, D)
+        v_exp = self._expand_kv(v)        # (B, H, T, D)
+
+        # Reshape gates for FLA: (B, T, d_kv) → (B, H, T, D) → reshape for kernel
+        erase = rearrange(erase_gate, "b t (h d) -> b h t d", h=H_kv)
+        write = rearrange(write_gate, "b t (h d) -> b h t d", h=H_kv)
+        alpha = rearrange(decay, "b t (h d) -> b h t d", h=H_kv)
+        
+        erase_exp = self._expand_kv(erase)  # (B, H, T, D)
+        write_exp = self._expand_kv(write)  # (B, H, T, D)
+        alpha_exp = self._expand_kv(alpha)  # (B, H, T, D)
+
+        try:
+            # Use FLA chunk kernel for all sequences
+            from fla.ops.gated_delta_rule import chunk_gated_delta_rule as fla_chunk_gdn
+            
+            # FLA expects: q,k = [B, T, H, D], v = [B, T, HV, V]
+            # g = [B, T, HV] (log-space forget gate, per-head)
+            # beta = [B, T, HV] (update gate, per-head, 0-1)
+            
+            # Transpose from (B, H, T, D) → (B, T, H, D) and ensure BF16
+            q_fla = q.transpose(1, 2).contiguous().to(torch.bfloat16)
+            k_fla = k_exp.transpose(1, 2).contiguous().to(torch.bfloat16)
+            
+            # Apply write gate to values, then transpose
+            v_gated = (write_exp * v_exp)                # (B, H, T, D)
+            v_fla = v_gated.transpose(1, 2).contiguous().to(torch.bfloat16)
+            
+            # g: per-head log-decay (MUST be float32)
+            g_per_head = torch.log(alpha_exp.float().clamp(min=1e-6)).mean(dim=-1)  # (B, H, T)
+            g_fla = g_per_head.transpose(1, 2).contiguous()  # (B, T, H) float32
+            
+            # beta: per-head erase strength (MUST be float32)
+            beta_per_head = erase_exp.float().mean(dim=-1)  # (B, H, T)
+            beta_fla = beta_per_head.transpose(1, 2).contiguous()  # (B, T, H) float32
+            
+            # Call FLA kernel outside autocast
+            with torch.amp.autocast(device_type="cuda", enabled=False):
+                output, _ = fla_chunk_gdn(
+                    q=q_fla,
+                    k=k_fla,
+                    v=v_fla,
+                    g=g_fla,
+                    beta=beta_fla,
+                    scale=1.0,
+                    use_qk_l2norm_in_kernel=False,
+                )
+            # output: (B, T, H, D) → transpose back to (B, H, T, D)
+            return output.transpose(1, 2).to(dtype)
+            
+        except (ImportError, RuntimeError) as e:
+            # Sequential recurrence fallback
+            return self._gdn2_sequential(q, k_exp, v_exp, erase_exp, write_exp, alpha_exp, attention_mask)
+
+    def _gdn2_sequential(
+        self,
+        q: torch.Tensor,      # (B, H, T, D)
+        k: torch.Tensor,      # (B, H, T, D)
+        v: torch.Tensor,      # (B, H, T, D)
+        erase: torch.Tensor,  # (B, H, T, D)
+        write: torch.Tensor,  # (B, H, T, D)
+        alpha: torch.Tensor,  # (B, H, T, D)
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Sequential fallback for GDN-2 (used when FLA kernel unavailable)."""
+        B, H, T, D = q.shape
+        device = q.device
+        dtype = q.dtype
+
+        state = torch.zeros(B, H, D, D, device=device, dtype=torch.float32)
+        outputs = []
+
+        for t in range(T):
+            if attention_mask is not None and t < attention_mask.shape[1]:
+                mask_t = attention_mask[:, t].view(B, 1, 1, 1).float()
+            else:
+                mask_t = 1.0
+
+            q_t = q[:, :, t, :]
+            k_t = k[:, :, t, :]
+            v_t = v[:, :, t, :]
+            b_t = erase[:, :, t, :]
+            w_t = write[:, :, t, :]
+            a_t = alpha[:, :, t, :]
+
+            state = state * a_t.unsqueeze(-1).float()
+            u_t = (b_t * k_t).float()
+            old_val = torch.einsum("bhd,bhde->bhe", u_t, state)
+            target = (w_t * v_t).float()
+            error = target - old_val
+            state = state + torch.einsum("bhd,bhe->bhde", u_t, error) * mask_t
+            o_t = torch.einsum("bhd,bhde->bhe", q_t.float(), state)
+            outputs.append(o_t)
+
+        output = torch.stack(outputs, dim=2).to(dtype)
+        return output
+
+    # ── Forward ─────────────────────────────────────────────────
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        loop_idx: Optional[int] = None,
+        lora_deltas: Optional[dict] = None,
+    ) -> torch.Tensor:
+        B, T, _ = hidden_states.shape
+
+        # === Projections (with optional LoRA) ===
+        def proj_with_lora(proj, x, name):
+            if lora_deltas and name in lora_deltas:
+                return F.linear(x, proj.weight + lora_deltas[name], proj.bias if proj.bias is not None else None)
+            return proj(x)
+
+        q_raw = proj_with_lora(self.q_proj, hidden_states, "q_proj")  # (B, T, H*D)
+        k_raw = proj_with_lora(self.k_proj, hidden_states, "k_proj")  # (B, T, H_kv*D)
+        v = proj_with_lora(self.v_proj, hidden_states, "v_proj")      # (B, T, H_kv*D)
+
+        # === Path 1: Softmax (standard RoPE attention) ===
+        q1 = rearrange(q_raw, "b t (h d) -> b h t d", h=self.num_heads)
+        k1 = rearrange(k_raw, "b t (h d) -> b h t d", h=self.num_kv_heads)
+        v1 = rearrange(v, "b t (h d) -> b h t d", h=self.num_kv_heads)
+
+        # RoDE depth signal
+        if loop_idx is not None and self._rode is not None:
+            cos_d, sin_d = self._rode(q1, loop_idx)
+            q1 = apply_depth_rotary_emb(q1.flatten(0, 1), cos_d, sin_d).view(B, self.num_heads, T, -1)
+            k1 = apply_depth_rotary_emb(k1.flatten(0, 1), cos_d, sin_d).view(B, self.num_kv_heads, T, -1)
+
+        # Q/K norm
+        q1 = self.q_norm(q1)
+        k1 = self.k_norm(k1)
+
+        # RoPE
+        cos, sin = self.rotary_emb(q1, position_ids)
+        k1_expanded = self._expand_kv(k1)
+        v1_expanded = self._expand_kv(v1)
+        q1_rope, k1_rope = apply_rotary_pos_emb(q1, k1_expanded, cos, sin)
+        # Determine if we can skip softmax (alpha=0 means pure GDN-2)
+        skip_softmax = (hasattr(self, '_alpha_override') and self._alpha_override is not None 
+                       and self._alpha_override == 0.0)
+
+        if not skip_softmax:
+            o_softmax = self._softmax_attention(q1_rope, k1_rope, v1_expanded, attention_mask)
+
+        if not self.use_dual_path:
+            # Pure softmax mode (no GDN-2)
+            attn_output = o_softmax
+        else:
+            # === Path 2: GDN-2 ===
+            # Apply short conv to Q,K for GDN-2 path
+            if self.q_conv is not None:
+                q2_raw = F.silu(self.q_conv(q_raw))
+                k2_raw = F.silu(self.k_conv(k_raw))
+            else:
+                q2_raw = q_raw
+                k2_raw = k_raw
+
+            q2 = rearrange(q2_raw, "b t (h d) -> b h t d", h=self.num_heads)
+            k2 = rearrange(k2_raw, "b t (h d) -> b h t d", h=self.num_kv_heads)
+            v2 = rearrange(v, "b t (h d) -> b h t d", h=self.num_kv_heads)
+
+            # L2 normalize Q,K for GDN-2 (best practice)
+            q2 = F.normalize(q2, p=2, dim=-1)
+            k2 = F.normalize(k2, p=2, dim=-1)
+
+            # Compute gates
+            erase_gate = torch.sigmoid(self.erase_proj(hidden_states))  # (B, T, d_kv)
+            write_gate = torch.sigmoid(self.write_proj(hidden_states))  # (B, T, d_kv)
+            
+            # Compute decay: α_t = exp(-exp(a) * softplus(W_f @ x + δ))
+            decay_input = self.decay_proj(hidden_states) + self.decay_delta  # (B, T, d_kv)
+            g_t = -torch.exp(self.decay_a) * F.softplus(decay_input)        # (B, T, d_kv)
+            decay = torch.exp(g_t)  # ∈ (0, 1]
+
+            o_gdn2 = self._gdn2_attention(q2, k2, v2, erase_gate, write_gate, decay, attention_mask)
+
+            # === Merge paths ===
+            # Use scheduled alpha if set, otherwise learnable
+            if hasattr(self, '_alpha_override') and self._alpha_override is not None:
+                alpha = self._alpha_override
+            else:
+                alpha = torch.sigmoid(self.path_mix_logit).view(1, self.num_heads, 1, 1)
+            
+            if skip_softmax:
+                attn_output = o_gdn2
+            elif isinstance(alpha, float) and alpha == 0.0:
+                attn_output = o_gdn2
+            else:
+                attn_output = alpha * o_softmax + (1.0 - alpha) * o_gdn2
+
+        # === Output gate (SiLU) ===
+        gate = F.silu(rearrange(
+            self.output_gate_proj(hidden_states), "b t (h d) -> b h t d", h=self.num_heads
+        ))
+        attn_output = gate * attn_output
+
+        # === Output projection ===
+        attn_output = rearrange(attn_output, "b h t d -> b t (h d)")
+        attn_output = proj_with_lora(self.o_proj, attn_output, "o_proj")
+        attn_output = self.dropout(attn_output)
+
+        return attn_output

modeling/gla_attention.py

@@ -0,0 +1,381 @@
+"""
+Dual-Path Gated Linear Attention (v2.0).
+
+Solves the "RoPE through feature map" problem by splitting into two parallel paths
+on the SAME full 64-dim Q/K:
+
+  Path 1 (RoPE Chunked Attention):
+      Q_rope, K_rope = RoPE(Q, K)
+      O_rope = ChunkedSoftmaxAttention(Q_rope, K_rope, V, chunk_size=C)
+      Complexity: O(N × C)  — local positional accuracy
+
+  Path 2 (Decay Linear Attention — Parallel Associative Scan):
+      Q_lin, K_lin = φ(Q), φ(K)          ← Q/K BEFORE RoPE
+      Full T×T decay matrix: M[t,i] = exp(cumsum(logγ)[t] - cumsum(logγ)[i])
+      O_decay = (Q·K^T ⊙ M) · V / normalizer  ← bidirectional
+      Complexity: O(T² × H) memory/compute — fast for T ≤ 2048, Blackwell-optimized
+
+  Merge: O = α · O_rope + (1-α) · O_decay   ← α learnable per head
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from typing import Optional, List
+from einops import rearrange
+
+from .utils import RMSNorm, YaRNRotaryEmbedding, apply_rotary_pos_emb, RoDE, apply_depth_rotary_emb
+
+
+class DualPathGLAAttention(nn.Module):
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        head_dim: int,
+        feature_map: str = "relu",
+        attention_dropout: float = 0.0,
+        layer_idx: int = 0,
+        chunk_size: int = 2048,
+        decay_init: List[float] = None,
+        use_data_dependent_decay: bool = False,
+        # RoPE params
+        max_position_embeddings: int = 131072,
+        rope_theta: float = 1000000.0,
+        rope_scaling_factor: float = 4.0,
+        rope_original_max_position: int = 32768,
+        rope_beta_fast: int = 32,
+        rope_beta_slow: int = 1,
+    ):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.num_heads = num_heads
+        self.num_kv_heads = num_kv_heads
+        self.head_dim = head_dim
+        self.num_kv_groups = num_heads // num_kv_heads
+        self.layer_idx = layer_idx
+        self.chunk_size = chunk_size
+        self.use_data_dependent_decay = use_data_dependent_decay
+
+        # Shared projections (transplanted from Gemma 4)
+        self.q_proj = nn.Linear(hidden_size, num_heads * head_dim, bias=False)
+        self.k_proj = nn.Linear(hidden_size, num_kv_heads * head_dim, bias=False)
+        self.v_proj = nn.Linear(hidden_size, num_kv_heads * head_dim, bias=False)
+        self.o_proj = nn.Linear(num_heads * head_dim, hidden_size, bias=False)
+
+        # Gate (new param — not from Gemma 4)
+        self.gate_proj = nn.Linear(hidden_size, num_heads * head_dim, bias=True)
+
+        # Q/K normalization
+        self.q_norm = RMSNorm(head_dim)
+        self.k_norm = RMSNorm(head_dim)
+
+        # RoPE for Path 1
+        self.rotary_emb = YaRNRotaryEmbedding(
+            dim=head_dim,
+            max_position_embeddings=max_position_embeddings,
+            base=rope_theta,
+            scaling_factor=rope_scaling_factor,
+            original_max_position=rope_original_max_position,
+            beta_fast=rope_beta_fast,
+            beta_slow=rope_beta_slow,
+        )
+
+        # Per-head decay γ for Path 2 — init from config
+        if decay_init is None:
+            decay_init = [0.99] * num_heads
+        init_logits = torch.tensor([math.log(g / (1.0 - g)) for g in decay_init])
+        self.log_decay = nn.Parameter(init_logits)  # sigmoid(log_decay) = γ
+
+        # Data-dependent decay projection (init zeros for stability)
+        if use_data_dependent_decay:
+            self.decay_proj = nn.Linear(hidden_size, num_heads, bias=False)
+            nn.init.zeros_(self.decay_proj.weight)
+
+        # Per-head stream mix α — init 0 → sigmoid(0) = 0.5
+        self.stream_mix_logit = nn.Parameter(torch.zeros(num_heads))
+
+        # Per-path RMSNorm before merge (prevents dead path from variance mismatch)
+        self.rope_path_norm = RMSNorm(head_dim)
+        self.decay_path_norm = RMSNorm(head_dim)
+
+        # Post-GLA RMSNorm (new param)
+        self.post_attn_norm = RMSNorm(head_dim)
+
+        # RoDE: Rotary Depth Embedding applied to Q/K (not hidden_states)
+        # Only created if depth_rotary_dims is set in config (injected from outside via loop_idx)
+        self._rode = None  # Will be set by HyperloopSegment if needed
+
+        self.feature_map_type = feature_map
+        self.dropout = nn.Dropout(attention_dropout)
+
+        # Backward weight: learnable per-layer, init sigmoid(-2.2) ≈ 0.1
+        # Model learns optimal forward/backward balance per layer
+        self.backward_logit = nn.Parameter(torch.tensor(-2.2))
+
+    def _feature_map(self, x: torch.Tensor) -> torch.Tensor:
+        if self.feature_map_type == "elu":
+            return F.elu(x) + 1.0
+        return F.relu(x) + 1e-4  # default: relu + eps
+
+    def _expand_kv(self, x: torch.Tensor) -> torch.Tensor:
+        if self.num_kv_groups == 1:
+            return x
+        x = x.unsqueeze(2).expand(-1, -1, self.num_kv_groups, -1, -1)
+        return x.reshape(x.shape[0], self.num_heads, x.shape[3], x.shape[4])
+
+    # ── Path 1: RoPE Chunked Attention ──────────────────────────
+
+    def _chunked_attention(
+        self,
+        q: torch.Tensor,   # (B, H, T, D) — already RoPE'd
+        k: torch.Tensor,   # (B, H, T, D)
+        v: torch.Tensor,   # (B, H, T, D)
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        B, H, T, D = q.shape
+        C = self.chunk_size
+
+        # Pad to multiple of chunk_size
+        pad_len = (C - T % C) % C
+        if pad_len > 0:
+            q = F.pad(q, (0, 0, 0, pad_len))
+            k = F.pad(k, (0, 0, 0, pad_len))
+            v = F.pad(v, (0, 0, 0, pad_len))
+            if attention_mask is not None:
+                attention_mask = F.pad(attention_mask, (0, pad_len), value=0)
+
+        T_padded = q.shape[2]
+        num_chunks = T_padded // C
+
+        # Reshape into chunks: (B, H, num_chunks, C, D)
+        q_c = q.view(B, H, num_chunks, C, D)
+        k_c = k.view(B, H, num_chunks, C, D)
+        v_c = v.view(B, H, num_chunks, C, D)
+
+        # Attention scores within each chunk: (B, H, num_chunks, C, C)
+        scale = 1.0 / math.sqrt(D)
+        scores = torch.einsum("bhnid,bhnjd->bhnij", q_c, k_c) * scale
+
+        # Mask padding within chunks
+        if attention_mask is not None:
+            chunk_mask = attention_mask.view(B, num_chunks, C)  # (B, nc, C)
+            # (B, 1, nc, C, 1) × (B, 1, nc, 1, C) → (B, 1, nc, C, C)
+            mask_2d = chunk_mask[:, None, :, :, None] * chunk_mask[:, None, :, None, :]
+            scores = scores.masked_fill(mask_2d == 0, float("-inf"))
+
+        attn_weights = F.softmax(scores, dim=-1)
+        attn_weights = torch.nan_to_num(attn_weights, nan=0.0)
+        attn_weights = self.dropout(attn_weights)
+
+        # Weighted sum: (B, H, num_chunks, C, D)
+        out = torch.einsum("bhnij,bhnjd->bhnid", attn_weights, v_c)
+
+        # Reshape back: (B, H, T_padded, D) → trim to (B, H, T, D)
+        out = out.view(B, H, T_padded, D)[:, :, :T, :]
+        return out
+
+    # ── Path 2: Decay Linear Attention (Bidirectional) ──────────
+
+    def _decay_attention_one_direction(
+        self,
+        q: torch.Tensor,   # (B, H, T, D) — feature-mapped, NOT RoPE'd
+        k: torch.Tensor,
+        v: torch.Tensor,
+        gamma: torch.Tensor,  # (B, H) or (B, T, H) for data-dependent
+        attention_mask: Optional[torch.Tensor] = None,
+        reverse: bool = False,
+    ) -> torch.Tensor:
+        """
+        Parallel Associative Scan implementation of Decay Linear Attention.
+        Complexity: O(T^2) memory/compute — extremely fast for T <= 2048.
+        Bypasses the slow O(T) Python loop and avoids Blackwell Triton hangs.
+        """
+        B, H, T, D = q.shape
+        D_v = v.shape[-1]
+
+        if reverse:
+            q = q.flip(dims=[2])
+            k = k.flip(dims=[2])
+            v = v.flip(dims=[2])
+            if attention_mask is not None:
+                attention_mask = attention_mask.flip(dims=[1])
+
+        # FP32 for numerical stability
+        orig_dtype = q.dtype
+        q, k, v = q.float(), k.float(), v.float()
+        
+        # Gamma handling: (H,) or (B, T, H)
+        if gamma.dim() == 1:
+            # Broadcast (H,) to (B, T, H)
+            g = gamma.view(1, 1, H).expand(B, T, H)
+        elif gamma.dim() == 2:
+            # Broadcast (B, H) to (B, T, H)
+            g = gamma.unsqueeze(1).expand(B, T, H)
+        else:
+            g = gamma # (B, T, H)
+
+        # 1. Compute cumulative decay logs: log_g = log(gamma)
+        log_g = torch.log(g.clamp(min=1e-8))
+        
+        # 2. Compute Decay Matrix M where M[t, i] = prod_{j=i+1}^t gamma_j
+        # M[t, i] = exp(cumsum(log_g)[t] - cumsum(log_g)[i])
+        # We use a T x T matrix approach for maximum speed at small T
+        
+        # cum_log_g: (B, T, H)
+        cum_log_g = torch.cumsum(log_g, dim=1) 
+        
+        # (B, H, T, 1)
+        clg = cum_log_g.transpose(1, 2).unsqueeze(-1)
+        # diff[t, i] = cum_log_g[t] - cum_log_g[i] -> (B, H, T, T)
+        diff = clg - clg.transpose(-1, -2)
+        
+        decay_matrix = torch.exp(diff)
+        
+        # Causal mask: only i <= t
+        m_idx = torch.arange(T, device=q.device)
+        mask = (m_idx.view(-1, 1) >= m_idx.view(1, -1)).to(q.dtype)
+        decay_matrix = decay_matrix * mask.view(1, 1, T, T)
+        
+        if attention_mask is not None:
+            # decay_matrix[t, i] should be 0 if token i is masked
+            decay_matrix = decay_matrix * attention_mask.view(B, 1, 1, T)
+
+        # 3. Compute output: O_t = sum_{i=0}^t M[t, i] * (Q_t * K_i^T) * V_i
+        
+        # dot_qk: (B, H, T, T)
+        dot_qk = torch.einsum("bhtd,bhid->bhti", q, k)
+        # weights: (B, H, T, T)
+        weights = dot_qk * decay_matrix
+        
+        # normalizer: sum of decay * K
+        # z_t = sum_{i=0}^t M[t, i] * K_i
+        z_scores = torch.einsum("bhti,bhid->bhtd", decay_matrix, k)
+        norm = torch.einsum("bhtd,bhtd->bht", q, z_scores).unsqueeze(-1).clamp(min=1e-6)
+        
+        # out: (B, H, T, D_v)
+        out = torch.einsum("bhti,bhid->bhtd", weights, v)
+        out = (out / norm).to(orig_dtype)
+
+        if reverse:
+            out = out.flip(dims=[2])
+        return out
+
+    @property
+    def backward_weight(self):
+        return torch.sigmoid(self.backward_logit)
+
+    def _decay_attention(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        gamma: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        backward_weight: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        bw = backward_weight if backward_weight is not None else self.backward_weight
+        fwd = self._decay_attention_one_direction(q, k, v, gamma, attention_mask, reverse=False)
+        bwd = self._decay_attention_one_direction(q, k, v, gamma, attention_mask, reverse=True)
+        return (1.0 - bw) * fwd + bw * bwd
+
+    # ── Forward ─────────────────────────────────────────────────
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        backward_weight: Optional[torch.Tensor] = None,
+        stream_mix_alpha: Optional[torch.Tensor] = None,
+        loop_idx: Optional[int] = None,
+        lora_deltas: Optional[dict] = None,
+    ) -> torch.Tensor:
+        B, T, _ = hidden_states.shape
+
+        # === Shared projections (with optional per-loop LoRA) ===
+        if lora_deltas and "q_proj" in lora_deltas:
+            q = F.linear(hidden_states, self.q_proj.weight + lora_deltas["q_proj"])
+        else:
+            q = self.q_proj(hidden_states)
+        q = rearrange(q, "b t (h d) -> b h t d", h=self.num_heads)
+
+        if lora_deltas and "k_proj" in lora_deltas:
+            k = F.linear(hidden_states, self.k_proj.weight + lora_deltas["k_proj"])
+        else:
+            k = self.k_proj(hidden_states)
+        k = rearrange(k, "b t (h d) -> b h t d", h=self.num_kv_heads)
+
+        if lora_deltas and "v_proj" in lora_deltas:
+            v = F.linear(hidden_states, self.v_proj.weight + lora_deltas["v_proj"])
+        else:
+            v = self.v_proj(hidden_states)
+        v = rearrange(v, "b t (h d) -> b h t d", h=self.num_kv_heads)
+
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        # === RoDE: inject depth signal into Q and K (correct location vs hidden_states) ===
+        # Q/K are in (B, H, T, D). RoDE xoay 8 dims đầu của head_dim.
+        if loop_idx is not None and self._rode is not None:
+            cos_d, sin_d = self._rode(q, loop_idx)   # (1,1,1, depth_dims)
+            q = apply_depth_rotary_emb(q.flatten(0,1), cos_d, sin_d).view(B, self.num_heads, T, -1)
+            k = apply_depth_rotary_emb(k.flatten(0,1), cos_d, sin_d).view(B, self.num_kv_heads, T, -1)
+
+        # Expand KV for GQA BEFORE both paths
+        k_expanded = self._expand_kv(k)
+        v_expanded = self._expand_kv(v)
+
+        # === Path 1: RoPE Chunked Attention ===
+        cos, sin = self.rotary_emb(q, position_ids)
+        q_rope, k_rope = apply_rotary_pos_emb(q, self._expand_kv(k), cos, sin)
+        o_rope = self._chunked_attention(q_rope, k_rope, v_expanded, attention_mask)
+
+        # === Path 2: Decay Linear Attention (on original Q/K, no RoPE) ===
+        q_lin = self._feature_map(q)
+        k_lin = self._feature_map(k_expanded)
+
+        # Compute gamma
+        gamma = torch.sigmoid(self.log_decay)  # (H,)
+        if self.use_data_dependent_decay and hasattr(self, 'decay_proj'):
+            # gamma_t = sigmoid(base_logit + data_term)
+            data_term = self.decay_proj(hidden_states)  # (B, T, H)
+            gamma = torch.sigmoid(
+                self.log_decay.view(1, 1, -1) + data_term
+            )  # (B, T, H) — will be handled in scan
+
+        o_decay = self._decay_attention(q_lin, k_lin, v_expanded, gamma, attention_mask, backward_weight=backward_weight)
+
+        # === Normalize each path before merge (prevents dead path) ===
+        o_rope = self.rope_path_norm(o_rope)
+        o_decay = self.decay_path_norm(o_decay)
+
+        # === Merge: α · O_rope + (1-α) · O_decay ===
+        alpha = stream_mix_alpha if stream_mix_alpha is not None else torch.sigmoid(self.stream_mix_logit)
+        alpha = alpha.view(1, self.num_heads, 1, 1)
+        attn_output = alpha * o_rope + (1.0 - alpha) * o_decay
+
+        # Post-GLA RMSNorm (final stabilization)
+        attn_output = self.post_attn_norm(attn_output)
+
+        # Gate
+        gate = F.silu(
+            rearrange(self.gate_proj(hidden_states), "b t (h d) -> b h t d", h=self.num_heads)
+        )
+        attn_output = gate * attn_output
+
+        # Store gate for weighted pooling (detached)
+        # gate: (B, H, T, D) → rearrange to (B, T, H*D) to match hidden_size
+        self._last_gate = rearrange(gate, "b h t d -> b t (h d)").detach()
+
+        # Output projection (with optional per-loop LoRA)
+        attn_output = rearrange(attn_output, "b h t d -> b t (h d)")
+        if lora_deltas and "o_proj" in lora_deltas:
+            attn_output = F.linear(attn_output, self.o_proj.weight + lora_deltas["o_proj"])
+        else:
+            attn_output = self.o_proj(attn_output)
+        attn_output = self.dropout(attn_output)
+
+        return attn_output

modeling/hybrid_layer.py

@@ -0,0 +1,130 @@
+"""
+DeepX v0.7 Layer: GDN-2 Attention + SwiGLU MLP with pre-norm residual.
+
+Supports both narrow (8h) and wide (16h) configurations.
+"""
+
+import torch
+import torch.nn as nn
+from typing import Optional
+
+from .gdn2_attention import GatedDeltaNet2Attention
+from .utils import RMSNorm, SwiGLUMLP
+
+
+class DeepXLayer(nn.Module):
+    """Gated DeltaNet-2 Attention + SwiGLU MLP with pre-norm residual."""
+
+    def __init__(
+        self,
+        hidden_size: int,
+        num_heads: int,
+        num_kv_heads: int,
+        head_dim: int,
+        intermediate_size: int,
+        config,
+        layer_idx: int = 0,
+    ):
+        super().__init__()
+        self.input_norm = RMSNorm(hidden_size, config.rms_norm_eps)
+        self.self_attn = GatedDeltaNet2Attention(
+            hidden_size=hidden_size,
+            num_heads=num_heads,
+            num_kv_heads=num_kv_heads,
+            head_dim=head_dim,
+            chunk_size=config.chunk_size,
+            attention_dropout=config.attention_dropout,
+            layer_idx=layer_idx,
+            use_dual_path=config.use_dual_path,
+            softmax_init_alpha=config.softmax_init_alpha,
+            use_short_conv=config.use_short_conv,
+            conv_kernel_size=config.conv_kernel_size,
+            max_position_embeddings=config.max_position_embeddings,
+            rope_theta=config.rope_theta,
+            rope_scaling_factor=config.rope_scaling_factor,
+            rope_original_max_position=config.rope_original_max_position,
+            rope_beta_fast=config.rope_beta_fast,
+            rope_beta_slow=config.rope_beta_slow,
+        )
+        self.post_attn_norm = RMSNorm(hidden_size, config.rms_norm_eps)
+        self.mlp = SwiGLUMLP(hidden_size, intermediate_size)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        loop_idx: Optional[int] = None,
+        lora_deltas: Optional[dict] = None,
+    ) -> torch.Tensor:
+        # Pre-norm + Attention + Residual
+        residual = x
+        x = self.input_norm(x)
+        x = self.self_attn(
+            x,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            loop_idx=loop_idx,
+            lora_deltas=lora_deltas,
+        )
+        x = residual + x
+
+        # Pre-norm + MLP + Residual
+        residual = x
+        x = self.post_attn_norm(x)
+        x = self.mlp(x, lora_deltas=lora_deltas)
+        x = residual + x
+
+        return x
+
+
+def make_narrow_a_layer(config, layer_idx: int = 0) -> DeepXLayer:
+    """Create NarrowA layer (8h, 1kv, MLP 6144)."""
+    return DeepXLayer(
+        hidden_size=config.hidden_size,
+        num_heads=config.narrow_a_heads,
+        num_kv_heads=config.narrow_a_kv_heads,
+        head_dim=config.narrow_a_head_dim,
+        intermediate_size=config.narrow_a_intermediate,
+        config=config,
+        layer_idx=layer_idx,
+    )
+
+
+def make_narrow_b_layer(config, layer_idx: int = 0) -> DeepXLayer:
+    """Create NarrowB layer (8h, 1kv, MLP 12288)."""
+    return DeepXLayer(
+        hidden_size=config.hidden_size,
+        num_heads=config.narrow_b_heads,
+        num_kv_heads=config.narrow_b_kv_heads,
+        head_dim=config.narrow_b_head_dim,
+        intermediate_size=config.narrow_b_intermediate,
+        config=config,
+        layer_idx=layer_idx,
+    )
+
+
+def make_wide_a_layer(config, layer_idx: int = 0) -> DeepXLayer:
+    """Create WideA layer (16h, 2kv, MLP 6144)."""
+    return DeepXLayer(
+        hidden_size=config.hidden_size,
+        num_heads=config.wide_a_heads,
+        num_kv_heads=config.wide_a_kv_heads,
+        head_dim=config.wide_a_head_dim,
+        intermediate_size=config.wide_a_intermediate,
+        config=config,
+        layer_idx=layer_idx,
+    )
+
+
+def make_wide_b_layer(config, layer_idx: int = 0) -> DeepXLayer:
+    """Create WideB layer (16h, 2kv, MLP 12288)."""
+    return DeepXLayer(
+        hidden_size=config.hidden_size,
+        num_heads=config.wide_b_heads,
+        num_kv_heads=config.wide_b_kv_heads,
+        head_dim=config.wide_b_head_dim,
+        intermediate_size=config.wide_b_intermediate,
+        config=config,
+        layer_idx=layer_idx,
+    )

modeling/hyperloop.py

@@ -0,0 +1,152 @@
+"""
+Hyperloop Segment v0.7 — Per-Loop LoRA + RoDE for Gated DeltaNet-2.
+
+Each loop iteration = [1 Wide + 4 Narrow] layers, with:
+  1. RoDE: Rotary depth signal on Q/K inside attention
+  2. Per-loop LoRA on all projections (Q/K/V/O + MLP gate/up/down)
+  3. Stochastic depth for robustness
+
+Two phase types:
+  Phase1: WideA(16h, MLP6144) + NarrowA×4(8h, MLP6144)
+  Phase2: NarrowB×4(8h, MLP12288) + WideB(16h, MLP12288)
+"""
+
+import torch
+import torch.nn as nn
+from typing import Optional, Dict
+
+
+class PerLoopLoRA(nn.Module):
+    """Per-loop low-rank adaptation for all projections."""
+
+    def __init__(self, num_loops: int, proj_shapes: Dict[str, tuple], rank: int = 16):
+        super().__init__()
+        self.num_loops = num_loops
+        self.rank = rank
+        self.proj_names = list(proj_shapes.keys())
+
+        for name, (out_dim, in_dim) in proj_shapes.items():
+            a_tensors = nn.ParameterList([
+                nn.Parameter(torch.zeros(out_dim, rank))
+                for _ in range(num_loops)
+            ])
+            b_tensors = nn.ParameterList([
+                nn.Parameter(torch.zeros(rank, in_dim))
+                for _ in range(num_loops)
+            ])
+            setattr(self, f"lora_A_{name}", a_tensors)
+            setattr(self, f"lora_B_{name}", b_tensors)
+
+    def get_delta(self, proj_name: str, loop_idx: int) -> torch.Tensor:
+        A = getattr(self, f"lora_A_{proj_name}")[loop_idx]
+        B = getattr(self, f"lora_B_{proj_name}")[loop_idx]
+        return A @ B
+
+
+class HyperloopPhase(nn.Module):
+    """
+    Multi-iteration loop with [Wide + Narrow×4] pattern per iteration.
+    
+    Each iteration:
+      1. Forward through shared_wide (1 pass)
+      2. Forward through shared_narrow × 4 (4 passes)
+    Total per iteration: 5 passes
+    """
+
+    def __init__(
+        self,
+        config,
+        shared_narrow: nn.Module,
+        shared_wide: nn.Module,
+        num_loops: int,
+        narrow_num_heads: int,
+        narrow_kv_heads: int,
+        narrow_head_dim: int,
+        narrow_intermediate: int,
+        wide_num_heads: int,
+        wide_kv_heads: int,
+        wide_head_dim: int,
+        wide_intermediate: int,
+        wide_first: bool = True,  # True: [Wide, Narrow×4], False: [Narrow×4, Wide]
+    ):
+        super().__init__()
+        self.shared_narrow = shared_narrow
+        self.shared_wide = shared_wide
+        self.num_loops = num_loops
+        self.drop_path_rate = config.drop_path_rate
+        self.wide_first = wide_first
+        H = config.hidden_size
+
+        # Per-loop LoRA for narrow layers (applied 4× per iteration)
+        narrow_proj_shapes = {
+            "q_proj": (narrow_num_heads * narrow_head_dim, H),
+            "k_proj": (narrow_kv_heads * narrow_head_dim, H),
+            "v_proj": (narrow_kv_heads * narrow_head_dim, H),
+            "o_proj": (H, narrow_num_heads * narrow_head_dim),
+            "gate_proj": (narrow_intermediate, H),
+            "up_proj": (narrow_intermediate, H),
+            "down_proj": (H, narrow_intermediate),
+        }
+        # Per-loop LoRA for wide layers (applied 1× per iteration)
+        wide_proj_shapes = {
+            "q_proj": (wide_num_heads * wide_head_dim, H),
+            "k_proj": (wide_kv_heads * wide_head_dim, H),
+            "v_proj": (wide_kv_heads * wide_head_dim, H),
+            "o_proj": (H, wide_num_heads * wide_head_dim),
+            "gate_proj": (wide_intermediate, H),
+            "up_proj": (wide_intermediate, H),
+            "down_proj": (H, wide_intermediate),
+        }
+
+        # LoRA for each iteration (narrow layers share LoRA within iteration)
+        self.narrow_lora = PerLoopLoRA(num_loops, narrow_proj_shapes, config.lora_rank)
+        self.wide_lora = PerLoopLoRA(num_loops, wide_proj_shapes, config.lora_rank)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+
+        for i in range(self.num_loops):
+            # Stochastic depth
+            if self.training and self.drop_path_rate > 0:
+                drop_prob = self.drop_path_rate * (i + 1) / self.num_loops
+                if torch.rand(1).item() < drop_prob:
+                    continue
+
+            # Get LoRA deltas for this iteration
+            narrow_deltas = {
+                name: self.narrow_lora.get_delta(name, i)
+                for name in self.narrow_lora.proj_names
+            }
+            wide_deltas = {
+                name: self.wide_lora.get_delta(name, i)
+                for name in self.wide_lora.proj_names
+            }
+
+            if self.wide_first:
+                # [Wide, Narrow×4]
+                hidden_states = self.shared_wide(
+                    hidden_states, attention_mask=attention_mask,
+                    position_ids=position_ids, loop_idx=i, lora_deltas=wide_deltas,
+                )
+                for _ in range(4):
+                    hidden_states = self.shared_narrow(
+                        hidden_states, attention_mask=attention_mask,
+                        position_ids=position_ids, loop_idx=i, lora_deltas=narrow_deltas,
+                    )
+            else:
+                # [Narrow×4, Wide]
+                for _ in range(4):
+                    hidden_states = self.shared_narrow(
+                        hidden_states, attention_mask=attention_mask,
+                        position_ids=position_ids, loop_idx=i, lora_deltas=narrow_deltas,
+                    )
+                hidden_states = self.shared_wide(
+                    hidden_states, attention_mask=attention_mask,
+                    position_ids=position_ids, loop_idx=i, lora_deltas=wide_deltas,
+                )
+
+        return hidden_states

modeling/mamba2_block.py

@@ -0,0 +1,244 @@
+"""
+Mamba2 Block wrapper for hybrid architecture.
+
+Mamba2 improves on Mamba1 with:
+- Multi-head SSM (similar to multi-head attention)
+- Larger effective state size
+- Hardware-efficient SSD (Structured State Space Duality) algorithm
+- 2-8x faster than Mamba1
+
+For embedding (bidirectional), we run Mamba2 in both directions
+and combine the outputs.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .utils import RMSNorm
+
+try:
+    from mamba_ssm import Mamba2 as Mamba2Core
+    MAMBA2_AVAILABLE = True
+except ImportError:
+    MAMBA2_AVAILABLE = False
+
+
+class Mamba2Fallback(nn.Module):
+    """
+    Pure PyTorch fallback when mamba_ssm is not installed.
+    Implements simplified SSM: y = SSM(Conv1d(Linear(x)))
+    Not as fast as CUDA kernel but functionally equivalent.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        d_state: int = 128,
+        d_conv: int = 4,
+        expand: int = 2,
+        headdim: int = 64,
+        ngroups: int = 8,
+    ):
+        super().__init__()
+        self.d_model = d_model
+        self.d_inner = d_model * expand
+        self.d_state = d_state
+        self.d_conv = d_conv
+        self.headdim = headdim
+        self.nheads = self.d_inner // headdim
+        self.ngroups = ngroups
+
+        assert self.d_inner % headdim == 0, (
+            f"d_inner ({self.d_inner}) must be divisible by headdim ({headdim})"
+        )
+        assert self.nheads % ngroups == 0, (
+            f"nheads ({self.nheads}) must be divisible by ngroups ({ngroups})"
+        )
+
+        # Input projection: x → (z, x_ssm)
+        self.in_proj = nn.Linear(d_model, self.d_inner * 2, bias=False)
+
+        # Conv1d for local context
+        self.conv1d = nn.Conv1d(
+            self.d_inner, self.d_inner,
+            kernel_size=d_conv, padding=d_conv - 1,
+            groups=self.d_inner, bias=True,
+        )
+
+        # SSM parameters
+        self.dt_proj = nn.Linear(self.d_inner, self.nheads, bias=True)
+        # A must be negative for SSM stability: A = -exp(A_log)
+        # Init A_log ~ log(uniform(1, 16)) per Mamba paper
+        A_init = torch.log(torch.rand(self.nheads) * 15 + 1)  # log(U(1,16))
+        self.A_log = nn.Parameter(A_init)
+        self.D = nn.Parameter(torch.ones(self.nheads))
+
+        # B, C projections (input-dependent)
+        self.B_proj = nn.Linear(self.d_inner, self.ngroups * d_state, bias=False)
+        self.C_proj = nn.Linear(self.d_inner, self.ngroups * d_state, bias=False)
+
+        # Output projection
+        self.out_proj = nn.Linear(self.d_inner, d_model, bias=False)
+        self.norm = nn.LayerNorm(self.d_inner)
+
+    def _ssm_scan(self, x, dt, A, B, C, D):
+        """
+        Parallel Associative Scan implementation of SSD (Structured State Space Duality).
+        Bypasses the slow O(T) Python loop and avoids Blackwell Triton hangs.
+        """
+        orig_dtype = x.dtype
+        x = x.float()
+        dt = dt.float()
+        A = A.float()
+        B = B.float()
+        C = C.float()
+
+        B_seq, T, d_inner = x.shape
+        nheads = self.nheads
+        headdim = self.headdim
+
+        x = x.view(B_seq, T, nheads, headdim)
+        dt = F.softplus(dt)  # (B, T, nheads)
+        A = -torch.exp(A)    # (nheads,)
+
+        # 1. Compute dA: A_bar[t] = dt[t] * A
+        log_dA = dt.unsqueeze(-1) * A.view(1, 1, nheads, 1) # (B, T, H, 1)
+        
+        # 2. Compute cumulative product of dA using cumsum of logs
+        # M[t, i] = prod_{j=i+1}^t exp(log_dA[j]) = exp(cumsum(log_dA)[t] - cumsum(log_dA)[i])
+        cum_log_dA = torch.cumsum(log_dA, dim=1) # (B, T, H, 1)
+        
+        # clg: (B, H, T, 1)
+        clg = cum_log_dA.permute(0, 2, 1, 3) 
+        # diff[t, i] = cum_log_dA[t] - cum_log_dA[i] -> (B, H, T, T)
+        diff = clg - clg.transpose(-1, -2)
+        decay_matrix = torch.exp(diff)
+        
+        # Causal mask
+        m_idx = torch.arange(T, device=x.device)
+        mask = (m_idx.view(-1, 1) >= m_idx.view(1, -1)).to(x.dtype)
+        decay_matrix = decay_matrix * mask.view(1, 1, T, T)
+
+        # 3. Compute B_bar * x: (B, T, H, S) * (B, T, H, D) -> states
+        B = B.view(B_seq, T, self.ngroups, self.d_state)
+        heads_per_group = nheads // self.ngroups
+        B = B.repeat_interleave(heads_per_group, dim=2)  # (B, T, H, S)
+        
+        # dB_x: (B, T, H, S, D)
+        dB = dt.unsqueeze(-1) * B # (B, T, H, S)
+        bx = torch.einsum("bths,bthd->bthsd", dB, x) 
+        
+        # 4. Apply decay matrix to states: h_t = sum_{i=0}^t M[t, i] * (dB_i * x_i)
+        # bx_flat: (B, H, T, S*D)
+        bx_flat = bx.permute(0, 2, 1, 3, 4).reshape(B_seq, nheads, T, -1)
+        h_flat = torch.einsum("bhti,bhis->bhts", decay_matrix, bx_flat)
+        h = h_flat.view(B_seq, nheads, T, self.d_state, headdim)
+        
+        # 5. Output: y_t = C_t * h_t
+        C = C.view(B_seq, T, self.ngroups, self.d_state)
+        C = C.repeat_interleave(heads_per_group, dim=2) # (B, T, H, S)
+        
+        # h_perm: (B, T, H, S, D)
+        h_perm = h.permute(0, 2, 1, 3, 4)
+        # y: (B, T, H, D)
+        y = torch.einsum("bths,bthsd->bthd", C, h_perm)
+        
+        # Add D skip connection
+        y = y + D.view(1, 1, nheads, 1) * x
+        
+        return y.reshape(B_seq, T, d_inner).to(orig_dtype)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if not hasattr(self, '_warned') and x.shape[1] > 512:
+            import warnings
+            warnings.warn(
+                f"Mamba2Fallback: sequential SSM scan with seq_len={x.shape[1]} will be slow. "
+                f"Install mamba_ssm for CUDA-accelerated scan.",
+                stacklevel=2,
+            )
+            self._warned = True
+
+        B, T, D = x.shape
+
+        # Input projection
+        xz = self.in_proj(x)  # (B, T, 2*d_inner)
+        x_ssm, z = xz.chunk(2, dim=-1)
+
+        # Conv1d
+        x_conv = x_ssm.transpose(1, 2)  # (B, d_inner, T)
+        x_conv = self.conv1d(x_conv)[:, :, :T]  # trim padding
+        x_conv = x_conv.transpose(1, 2)  # (B, T, d_inner)
+        x_conv = F.silu(x_conv)
+
+        # SSM parameters from input
+        dt = self.dt_proj(x_conv)
+        B_param = self.B_proj(x_conv)
+        C_param = self.C_proj(x_conv)
+
+        # SSM scan
+        y = self._ssm_scan(x_conv, dt, self.A_log, B_param, C_param, self.D)
+
+        # Gate with z
+        y = y * F.silu(z)
+        y = self.out_proj(self.norm(y))
+
+        return y
+
+
+class Mamba2Block(nn.Module):
+    """
+    Bidirectional Mamba2 block for embedding.
+    Runs Mamba2 forward + backward, combines outputs.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        d_state: int = 128,
+        d_conv: int = 4,
+        expand: int = 2,
+        headdim: int = 64,
+        ngroups: int = 8,
+        bidirectional: bool = True,
+    ):
+        super().__init__()
+        self.bidirectional = bidirectional
+        
+        mamba_cls = Mamba2Core if MAMBA2_AVAILABLE else Mamba2Fallback
+        mamba_kwargs = dict(
+            d_model=d_model,
+            d_state=d_state,
+            d_conv=d_conv,
+            expand=expand,
+            headdim=headdim,
+            ngroups=ngroups,
+        )
+
+        self.forward_mamba = mamba_cls(**mamba_kwargs)
+
+        if bidirectional:
+            self.backward_mamba = mamba_cls(**mamba_kwargs)
+            self.merge_proj = nn.Linear(d_model * 2, d_model, bias=False)
+
+    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None) -> torch.Tensor:
+        mask = None
+        if attention_mask is not None:
+            mask = attention_mask.unsqueeze(-1).to(x.dtype)
+            x = x * mask
+
+        fwd_out = self.forward_mamba(x)
+
+        if self.bidirectional:
+            x_flip = x.flip(dims=[1])
+            if mask is not None:
+                x_flip = x_flip * mask.flip(dims=[1])
+            bwd_out = self.backward_mamba(x_flip).flip(dims=[1])
+            out = self.merge_proj(torch.cat([fwd_out, bwd_out], dim=-1))
+        else:
+            out = fwd_out
+
+        # Re-apply mask: SSM hidden state decay để lại artifact tại padding positions
+        if mask is not None:
+            out = out * mask
+
+        return out

modeling/pipeline.py

@@ -0,0 +1,199 @@
+"""
+DeepX v0.6 Full Pipeline.
+
+Combines:
+  1. Frozen Gemma 4 E2B token embedding (loaded from pretrained/gemma4_e2b_embed.pt)
+  2. Pure GLA Hyperloop backbone + ColBERT head (PureGLAEmbeddingModel)
+
+Outputs:
+  - encode()         → single vector (1536-d) for fast ANN retrieval
+  - encode_colbert() → token vectors (T × 128-d) for MaxSim reranking
+  - encode_multi()   → both single + token vectors in one forward pass
+
+Weight Init: ~90% of backbone can be copied from Gemma 4 E2B.
+"""
+
+import torch
+import torch.nn as nn
+import logging
+import dataclasses
+from typing import Optional, Tuple
+from pathlib import Path
+
+from config import HybridEmbeddingConfig
+from .embedding_model import DeepXEmbeddingModel
+
+logger = logging.getLogger(__name__)
+
+
+class DeepXPipeline(nn.Module):
+    """
+    Full DeepX v0.6 embedding pipeline.
+
+    Token embedding is frozen and loaded from a pre-extracted file.
+    Only the backbone (PureGLAEmbeddingModel) is trained.
+    """
+
+    def __init__(
+        self,
+        config: HybridEmbeddingConfig,
+        embed_path: str = "pretrained/gemma4_e2b_embed.pt",
+    ):
+        super().__init__()
+        self.config = config
+
+        # --- Frozen Token Embedding (Gemma 4 E2B) ---
+        embed_path = Path(embed_path)
+        if not embed_path.exists():
+            raise FileNotFoundError(
+                f"Token embedding not found at '{embed_path}'.\n"
+                f"Please run: python scripts/extract_gemma_embedding.py"
+            )
+
+        logger.info(f"Loading frozen token embedding from {embed_path} ...")
+        weight = torch.load(embed_path, weights_only=True)
+
+        assert weight.shape == (config.vocab_size, config.hidden_size), (
+            f"Embedding shape mismatch: expected ({config.vocab_size}, {config.hidden_size}), "
+            f"got {tuple(weight.shape)}. Check config.hidden_size matches E2B."
+        )
+
+        self.token_embedding = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.token_embedding.weight.data = weight.to(config.torch_dtype)
+        self.token_embedding.requires_grad_(False)
+        logger.info(f"Token embedding frozen. Shape: {weight.shape}, dtype: {config.torch_dtype}")
+
+        # --- Trainable Backbone ---
+        self.backbone = DeepXEmbeddingModel(config)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        normalize: bool = True,
+        truncate_dim: Optional[int] = None,
+        return_colbert: bool = False,
+    ):
+        """
+        Full forward pass.
+        
+        Returns:
+            If return_colbert=False: single_embed (B, D)
+            If return_colbert=True: (single_embed (B, D), token_embeds (B, T, colbert_dim))
+        """
+        with torch.no_grad():
+            hidden_states = self.token_embedding(input_ids)
+
+        return self.backbone(
+            hidden_states,
+            attention_mask=attention_mask,
+            normalize=normalize,
+            truncate_dim=truncate_dim,
+            return_colbert=return_colbert,
+        )
+
+    def encode(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        truncate_dim: Optional[int] = None,
+    ) -> torch.Tensor:
+        """Single vector encoding for fast ANN retrieval."""
+        with torch.no_grad():
+            return self.forward(input_ids, attention_mask, normalize=True, truncate_dim=truncate_dim)
+
+    def encode_colbert(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        ColBERT encoding — returns both single vector and token vectors.
+        
+        Returns:
+            single_embed: (B, 1536) for coarse retrieval
+            token_embeds: (B, T, 128) for MaxSim reranking
+        """
+        with torch.no_grad():
+            return self.forward(input_ids, attention_mask, normalize=True, return_colbert=True)
+
+    def encode_multi(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        truncate_dim: Optional[int] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """Alias for encode_colbert with optional truncation on single vector."""
+        with torch.no_grad():
+            hidden_states = self.token_embedding(input_ids)
+        return self.backbone(
+            hidden_states,
+            attention_mask=attention_mask,
+            normalize=True,
+            truncate_dim=truncate_dim,
+            return_colbert=True,
+        )
+
+    def freeze_embedder(self):
+        """Ensure token embedding stays frozen."""
+        self.token_embedding.requires_grad_(False)
+
+    def unfreeze_embedder(self):
+        """Unfreeze token embedding for fine-tuning (use with very small LR ~1e-6)."""
+        self.token_embedding.requires_grad_(True)
+        logger.warning("Token embedding UNFROZEN. Use lr ~1e-6.")
+
+    def count_parameters(self) -> dict:
+        embed_params = self.token_embedding.weight.numel()
+        backbone_counts = self.backbone.count_parameters()
+        return {
+            "embedding_frozen": embed_params,
+            "backbone_trainable": backbone_counts["trainable"],
+            "backbone_total": backbone_counts["backbone_total"],
+            "grand_total": embed_params + backbone_counts["backbone_total"],
+        }
+
+    def save_backbone(self, path: str) -> None:
+        """Save only the trained backbone weights."""
+        out = Path(path)
+        out.parent.mkdir(parents=True, exist_ok=True)
+        torch.save({
+            "state_dict": self.backbone.state_dict(),
+            "config": dataclasses.asdict(self.config),
+            "version": "0.7",
+            "architecture": "gdn2_hyperloop_colbert",
+            "embed_source": "gemma4_e2b",
+            "hidden_size": self.config.hidden_size,
+            "vocab_size": self.config.vocab_size,
+            "colbert_dim": self.config.colbert_dim,
+        }, out)
+        size_mb = out.stat().st_size / 1024 / 1024
+        logger.info(f"Backbone saved to {out} ({size_mb:.1f} MB)")
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        config: HybridEmbeddingConfig,
+        embed_path: str,
+        backbone_path: str,
+    ) -> "DeepXPipeline":
+        """Load pipeline from 2 .pt files for deployment."""
+        backbone_path = Path(backbone_path)
+        if not backbone_path.exists():
+            raise FileNotFoundError(f"Backbone weights not found at '{backbone_path}'.")
+
+        pipeline = cls(config, embed_path=embed_path)
+
+        logger.info(f"Loading backbone from {backbone_path} ...")
+        checkpoint = torch.load(backbone_path, weights_only=True, map_location="cpu")
+
+        if isinstance(checkpoint, dict) and "state_dict" in checkpoint:
+            state_dict = checkpoint["state_dict"]
+        else:
+            state_dict = checkpoint
+
+        pipeline.backbone.load_state_dict(state_dict)
+        pipeline.backbone.to(config.torch_dtype)
+        logger.info("Backbone loaded successfully.")
+
+        return pipeline

modeling/utils.py

@@ -0,0 +1,224 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from typing import Optional, Tuple
+
+
+class RMSNorm(nn.Module):
+    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: torch.Tensor) -> torch.Tensor:
+        input_dtype = x.dtype
+        x = x.float()
+        variance = x.pow(2).mean(-1, keepdim=True)
+        x = x * torch.rsqrt(variance + self.eps)
+        return (self.weight.float() * x).to(input_dtype)
+
+
+class SwiGLUMLP(nn.Module):
+    def __init__(self, hidden_size: int, intermediate_size: int):
+        super().__init__()
+        self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False)
+        self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False)
+
+    def forward(self, x: torch.Tensor, lora_deltas: dict = None) -> torch.Tensor:
+        if lora_deltas:
+            gate = F.linear(x, self.gate_proj.weight + lora_deltas.get("gate_proj", 0))
+            up = F.linear(x, self.up_proj.weight + lora_deltas.get("up_proj", 0))
+            down_weight = self.down_proj.weight + lora_deltas.get("down_proj", 0)
+            return F.linear(F.silu(gate) * up, down_weight)
+        return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+def _yarn_find_correction_dim(
+    num_rotations: int, dim: int, base: float = 10000.0, max_position: int = 2048
+) -> float:
+    """Find correction dimension for YaRN interpolation."""
+    return (dim * math.log(max_position / (num_rotations * 2 * math.pi))) / (2 * math.log(base))
+
+
+def _yarn_find_correction_range(
+    low_rot: int, high_rot: int, dim: int, base: float = 10000.0, max_position: int = 2048
+) -> Tuple[int, int]:
+    """Find the range of dimensions to apply YaRN correction."""
+    low = math.floor(_yarn_find_correction_dim(low_rot, dim, base, max_position))
+    high = math.ceil(_yarn_find_correction_dim(high_rot, dim, base, max_position))
+    return max(low, 0), min(high, dim - 1)
+
+
+def _yarn_linear_ramp_mask(low: int, high: int, dim: int, dtype: torch.dtype) -> torch.Tensor:
+    """Create linear ramp mask for smooth interpolation between dimensions."""
+    if low == high:
+        high += 0.001
+    linear_func = (torch.arange(dim, dtype=dtype) - low) / (high - low)
+    return linear_func.clamp(0, 1)
+
+
+class YaRNRotaryEmbedding(nn.Module):
+    """
+    RoPE with YaRN (Yet another RoPE extensioN) scaling.
+    
+    Same approach as Gemma 4 E2B:
+    - Base theta = 1,000,000
+    - YaRN scaling for context extension to 128K
+    - Splits dimensions into 3 regions:
+      1. Low freq dims: apply NTK-aware interpolation
+      2. Medium freq dims: smooth ramp between interpolation and extrapolation  
+      3. High freq dims: no scaling (extrapolation)
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        max_position_embeddings: int = 131072,
+        base: float = 1000000.0,
+        scaling_factor: float = 4.0,
+        original_max_position: int = 32768,
+        beta_fast: int = 32,
+        beta_slow: int = 1,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.max_position_embeddings = max_position_embeddings
+        self.base = base
+        self.scaling_factor = scaling_factor
+        self.original_max_position = original_max_position
+        self.beta_fast = beta_fast
+        self.beta_slow = beta_slow
+
+        self._build_yarn_cache()
+
+    def _build_yarn_cache(self):
+        """Compute YaRN-adjusted inverse frequencies."""
+        dim = self.dim
+        # Standard RoPE inverse frequencies
+        inv_freq = 1.0 / (
+            self.base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)
+        )
+
+        # YaRN correction
+        low, high = _yarn_find_correction_range(
+            self.beta_slow, self.beta_fast, dim, self.base, self.original_max_position
+        )
+        inv_freq_mask = 1.0 - _yarn_linear_ramp_mask(low, high, dim // 2, torch.float32)
+
+        # Interpolated frequencies (for extending context)
+        inv_freq_interpolated = inv_freq / self.scaling_factor
+
+        # Blend: high freq dims keep original, low freq dims get interpolated
+        inv_freq_yarn = inv_freq_interpolated * (1 - inv_freq_mask) + inv_freq * inv_freq_mask
+
+        self.register_buffer("inv_freq", inv_freq_yarn, persistent=False)
+
+        # Attention scaling factor (magnitude correction)
+        self.attn_scale = 0.1 * math.log(self.scaling_factor) + 1.0
+
+    def forward(
+        self, x: torch.Tensor, position_ids: Optional[torch.Tensor] = None
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Args:
+            x: (B, H, T, D) — used only for device/dtype
+            position_ids: (B, T) or None (auto-generate 0..T-1)
+        Returns:
+            cos, sin: (B, T, D) in same dtype as x
+        """
+        B, H, T, D = x.shape
+
+        if position_ids is None:
+            position_ids = torch.arange(T, device=x.device).unsqueeze(0).expand(B, -1)
+
+        # Compute in float32 for precision, cast output to match x
+        inv_freq = self.inv_freq.to(device=x.device, dtype=torch.float32)
+        freqs = position_ids.unsqueeze(-1).float() * inv_freq.unsqueeze(0).unsqueeze(0)
+
+        emb = torch.cat([freqs, freqs], dim=-1)
+
+        cos = (emb.cos() * self.attn_scale).to(x.dtype)
+        sin = (emb.sin() * self.attn_scale).to(x.dtype)
+
+        return cos, sin
+
+
+def apply_rotary_pos_emb(
+    q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Apply RoPE rotation to Q and K.
+    
+    Args:
+        q: (B, H, T, D)
+        k: (B, H_kv, T, D)
+        cos: (B, T, D)
+        sin: (B, T, D)
+    Returns:
+        q_rotated, k_rotated: same shapes
+    """
+    # (B, T, D) → (B, 1, T, D) for broadcasting with heads
+    cos = cos.unsqueeze(1)
+    sin = sin.unsqueeze(1)
+
+    q_rotated = (q * cos) + (_rotate_half(q) * sin)
+    k_rotated = (k * cos) + (_rotate_half(k) * sin)
+
+    return q_rotated, k_rotated
+
+
+def _rotate_half(x: torch.Tensor) -> torch.Tensor:
+    """Rotate half the hidden dims: [x1, x2] → [-x2, x1]"""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_depth_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+    """
+    Apply depth rotation to a subset of dimensions.
+    Args:
+        x: (B, T, D)
+        cos, sin: (1, 1, d_rode) where d_rode is the number of rotated dimensions.
+    """
+    d_rode = cos.shape[-1]
+    x_rode = x[..., :d_rode]
+    x_rest = x[..., d_rode:]
+    
+    # Apply rotation to the first d_rode dimensions
+    x_rotated = (x_rode * cos) + (_rotate_half(x_rode) * sin)
+    
+    return torch.cat([x_rotated, x_rest], dim=-1)
+
+
+class RoDE(nn.Module):
+    """
+    Rotary Depth Embedding (RoDE).
+    Provides a depth signal for shared weights in Hyperloop.
+    """
+    def __init__(self, dim: int, num_loops: int, base: float = 10000.0):
+        super().__init__()
+        self.dim = dim
+        self.num_loops = num_loops
+        self.base = base
+        
+        # Pre-compute sin/cos for each loop index
+        inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
+        
+        # (num_loops, dim // 2)
+        loop_ids = torch.arange(num_loops).float()
+        freqs = loop_ids.unsqueeze(-1) * inv_freq.unsqueeze(0)
+        
+        # (num_loops, dim)
+        emb = torch.cat([freqs, freqs], dim=-1)
+        self.register_buffer("cos", emb.cos(), persistent=False)
+        self.register_buffer("sin", emb.sin(), persistent=False)
+
+    def forward(self, x: torch.Tensor, loop_idx: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        # Pick the pre-computed sin/cos for the current loop index
+        # Shape: (1, 1, dim) for broadcasting
+        cos = self.cos[loop_idx].view(1, 1, -1).to(x.dtype)
+        sin = self.sin[loop_idx].view(1, 1, -1).to(x.dtype)
+        return cos, sin