scrapegoat-lora/modeling_scrapegoat.py

import math
import logging
from typing import List, Optional, Tuple, Union

import torch
import torch.nn as nn
import torch.nn.functional as F
from transformers.modeling_utils import PreTrainedModel
from transformers.modeling_outputs import (
    BaseModelOutputWithPast,
    CausalLMOutputWithPast,
)

from configuration_scrapegoat import ScrapeGoatConfig
from dspark_components import DSparkAttention, DSparkMarkovHead, DSparkConfidenceHead

logger = logging.getLogger(__name__)


def activation_fn(x, act_type="silu"):
    if act_type == "silu":
        return F.silu(x)
    elif act_type == "situ":
        return x * torch.sigmoid(x) * torch.tanh(x)
    elif act_type == "situ_simple":
        return x * torch.sigmoid(1.7 * x)
    else:
        return F.silu(x)


class ScrapeGoatRMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states):
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return self.weight * hidden_states.to(input_dtype)


class ScrapeGoatRotaryEmbedding(nn.Module):
    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
        super().__init__()
        self.dim = dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base
        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

    def forward(self, x, seq_len=None):
        if seq_len is None:
            seq_len = x.shape[-2]
        t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)
        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        emb = torch.cat((freqs, freqs), dim=-1)
        return emb.cos()[None, None, :, :], emb.sin()[None, None, :, :]


def rotate_half(x):
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None):
    if position_ids is None:
        cos = cos[:, :, : q.size(2), :]
        sin = sin[:, :, : q.size(2), :]
    elif position_ids.dim() == 2:
        cos = cos[:, :, position_ids[0], :]
        sin = sin[:, :, position_ids[0], :]
    else:
        cos = cos[:, :, position_ids, :]
        sin = sin[:, :, position_ids, :]
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed


class ScrapeGoatSourceAttention(nn.Module):
    """GQA attention for a single track with configurable head params."""

    def __init__(self, config, num_heads, num_kv_heads, head_dim):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.num_key_value_heads = num_kv_heads
        self.num_key_value_groups = num_heads // num_kv_heads
        self.max_position_embeddings = config.max_position_embeddings
        self.rope_theta = config.rope_theta
        self.is_causal = True

        self.q_proj = nn.Linear(self.hidden_size, num_heads * head_dim, bias=config.attention_bias)
        self.k_proj = nn.Linear(self.hidden_size, num_kv_heads * head_dim, bias=config.attention_bias)
        self.v_proj = nn.Linear(self.hidden_size, num_kv_heads * head_dim, bias=config.attention_bias)
        self.o_proj = nn.Linear(num_heads * head_dim, self.hidden_size, bias=config.attention_bias)
        self.rotary_emb = ScrapeGoatRotaryEmbedding(
            head_dim,
            max_position_embeddings=self.max_position_embeddings,
            base=self.rope_theta,
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        output_attentions: bool = False,
        use_cache: bool = False,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
        bsz, q_len, _ = hidden_states.size()

        query_states = self.q_proj(hidden_states)
        key_states = self.k_proj(hidden_states)
        value_states = self.v_proj(hidden_states)

        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
        key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)
        value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)

        kv_seq_len = key_states.shape[-2]
        if past_key_value is not None:
            kv_seq_len += past_key_value[0].shape[-2]

        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)

        if past_key_value is not None:
            key_states = torch.cat([past_key_value[0], key_states], dim=2)
            value_states = torch.cat([past_key_value[1], value_states], dim=2)

        present_key_value = (key_states, value_states) if use_cache else None

        key_states = key_states.repeat_interleave(self.num_key_value_groups, dim=1)
        value_states = value_states.repeat_interleave(self.num_key_value_groups, dim=1)

        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)

        if attention_mask is not None:
            attn_weights = attn_weights + attention_mask

        attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
        attn_output = torch.matmul(attn_weights, value_states)

        attn_output = attn_output.transpose(1, 2).contiguous()
        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
        attn_output = self.o_proj(attn_output)

        if not output_attentions:
            attn_weights = None

        return attn_output, attn_weights, present_key_value


class ScrapeGoatKDA(nn.Module):
    """Kimi Delta Attention — per-channel diagonal gating over Gated DeltaNet.

    Implements both recurrent (inference) and chunkwise (training) modes.
    Core equation (per-head):
        S_t = (I - β_t k_t k_t^T) Diag(α_t) S_{t-1} + β_t k_t v_t^T
        o_t = RMSNorm(S_t q_t) ⊙ sigmoid(W_g x_t)
    """

    def __init__(self, config, num_heads, head_dim):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.conv_kernel = config.kda_conv_kernel
        self.act_fn = lambda x: activation_fn(x, config.hidden_act)
        pad = self.conv_kernel // 2

        self.shortconv_q = nn.Conv1d(self.hidden_size, self.hidden_size, self.conv_kernel, padding=pad, groups=self.hidden_size, bias=False)
        self.shortconv_k = nn.Conv1d(self.hidden_size, self.hidden_size, self.conv_kernel, padding=pad, groups=self.hidden_size, bias=False)
        self.shortconv_v = nn.Conv1d(self.hidden_size, self.hidden_size, self.conv_kernel, padding=pad, groups=self.hidden_size, bias=False)

        self.q_proj = nn.Linear(self.hidden_size, num_heads * head_dim, bias=False)
        self.k_proj = nn.Linear(self.hidden_size, num_heads * head_dim, bias=False)
        self.v_proj = nn.Linear(self.hidden_size, num_heads * head_dim, bias=False)
        self.o_proj = nn.Linear(num_heads * head_dim, self.hidden_size, bias=False)

        self.alpha_down = nn.Linear(self.hidden_size, head_dim, bias=False)
        self.alpha_up = nn.Linear(head_dim, num_heads * head_dim, bias=False)

        self.beta_proj = nn.Linear(self.hidden_size, num_heads, bias=False)

        self.gate_proj = nn.Linear(self.hidden_size, num_heads * head_dim, bias=False)

        self.output_norm = ScrapeGoatRMSNorm(head_dim)

    def _apply_shortconv(self, x, conv):
        bsz, seq_len, h = x.shape
        x_t = x.transpose(1, 2)
        x_t = self.act_fn(x_t)
        x_t = conv(x_t)
        return x_t.transpose(1, 2)

    def _compute_alpha(self, x):
        bsz, sl, _ = x.shape
        h = self.alpha_down(x)
        h = self.act_fn(h)
        h = self.alpha_up(h)
        h = F.sigmoid(h)
        return h.view(bsz, sl, self.num_heads, self.head_dim)

    def _compute_beta(self, x):
        return F.sigmoid(self.beta_proj(x))

    def _forward_recurrent(self, q, k, v, alpha, beta, state=None):
        batch, nh, hd = q.shape
        if state is None:
            state = torch.zeros(batch, nh, hd, hd, device=q.device, dtype=q.dtype)

        k_u = k.unsqueeze(-1)
        v_u = v.unsqueeze(-1)
        q_u = q.unsqueeze(-1)

        alpha_diag = alpha.unsqueeze(-1)
        state_scaled = state * alpha_diag

        k_state = torch.matmul(k_u.transpose(-2, -1), state_scaled)
        beta_k = beta.view(batch, nh, 1, 1)
        state = state_scaled - beta_k * (k_u @ k_state)

        state = state + beta_k * (k_u @ v_u.transpose(-2, -1))

        out = (state @ q_u).squeeze(-1)
        return out, state

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        output_attentions: bool = False,
        use_cache: bool = False,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
        bsz, seq_len, h = hidden_states.shape

        q_sconv = self._apply_shortconv(hidden_states, self.shortconv_q)
        k_sconv = self._apply_shortconv(hidden_states, self.shortconv_k)
        v_sconv = self._apply_shortconv(hidden_states, self.shortconv_v)

        q = self.q_proj(q_sconv).view(bsz, seq_len, self.num_heads, self.head_dim)
        k = self.k_proj(k_sconv).view(bsz, seq_len, self.num_heads, self.head_dim)
        v = self.v_proj(v_sconv).view(bsz, seq_len, self.num_heads, self.head_dim)

        q = F.normalize(q, dim=-1)
        k = F.normalize(k, dim=-1)

        alpha = self._compute_alpha(hidden_states)
        beta = self._compute_beta(hidden_states)

        past_state = None
        if past_key_value is not None and past_key_value[0] is not None:
            past_state = past_key_value[0]

        if seq_len == 1 and past_state is not None:
            q_i = q[:, 0]
            k_i = k[:, 0]
            v_i = v[:, 0]
            a_i = alpha[:, 0]
            b_i = beta[:, 0]
            out_i, new_state = self._forward_recurrent(q_i, k_i, v_i, a_i, b_i, past_state)
            out_i = self.output_norm(out_i)
            gate = F.sigmoid(self.gate_proj(hidden_states[:, 0])).view(bsz, self.num_heads, self.head_dim)
            out_i = out_i * gate
            out = out_i.reshape(bsz, 1, -1)
            out = self.o_proj(out)
            present = (new_state,)
            return out, None, present

        outputs = []
        state = past_state
        for t in range(seq_len):
            q_i = q[:, t]
            k_i = k[:, t]
            v_i = v[:, t]
            a_i = alpha[:, t]
            b_i = beta[:, t]
            out_i, state = self._forward_recurrent(q_i, k_i, v_i, a_i, b_i, state)
            outputs.append(out_i)

        out = torch.stack(outputs, dim=1)
        out = self.output_norm(out)
        gate = F.sigmoid(self.gate_proj(hidden_states)).view(bsz, seq_len, self.num_heads, self.head_dim)
        out = out * gate
        out = out.reshape(bsz, seq_len, -1)
        out = self.o_proj(out)

        present = (state,) if use_cache else None

        return out, None, present


class ScrapeGoatMoMKDA(nn.Module):
    """Mixture-of-Memories KDA — multiple independent memory states with routing.

    Each state has its own K/V projections + alpha/beta/gate. Q projection shared.
    Router assigns each token to top-K states. Shared memory always active.
    """

    def __init__(self, config, num_heads, head_dim):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.num_memories = config.mom_num_memories
        self.active_memories = config.mom_active_memories
        self.shared_memory = config.mom_shared_memory
        self.act_fn = lambda x: activation_fn(x, config.hidden_act)
        conv_kernel = config.kda_conv_kernel
        pad = conv_kernel // 2

        self.shortconv_q = nn.Conv1d(self.hidden_size, self.hidden_size, conv_kernel, padding=pad, groups=self.hidden_size, bias=False)
        self.q_proj = nn.Linear(self.hidden_size, num_heads * head_dim, bias=False)
        self.o_proj = nn.Linear(num_heads * head_dim, self.hidden_size, bias=False)
        self.output_norm = ScrapeGoatRMSNorm(head_dim)

        self.shortconv_k = nn.ModuleList([
            nn.Conv1d(self.hidden_size, self.hidden_size, conv_kernel, padding=pad, groups=self.hidden_size, bias=False)
            for _ in range(self.num_memories)
        ])
        self.shortconv_v = nn.ModuleList([
            nn.Conv1d(self.hidden_size, self.hidden_size, conv_kernel, padding=pad, groups=self.hidden_size, bias=False)
            for _ in range(self.num_memories)
        ])
        self.k_proj = nn.ModuleList([
            nn.Linear(self.hidden_size, num_heads * head_dim, bias=False) for _ in range(self.num_memories)
        ])
        self.v_proj = nn.ModuleList([
            nn.Linear(self.hidden_size, num_heads * head_dim, bias=False) for _ in range(self.num_memories)
        ])
        self.alpha_down = nn.ModuleList([
            nn.Linear(self.hidden_size, head_dim, bias=False) for _ in range(self.num_memories)
        ])
        self.alpha_up = nn.ModuleList([
            nn.Linear(head_dim, num_heads * head_dim, bias=False) for _ in range(self.num_memories)
        ])
        self.beta_proj = nn.ModuleList([
            nn.Linear(self.hidden_size, num_heads, bias=False) for _ in range(self.num_memories)
        ])
        self.gate_proj = nn.ModuleList([
            nn.Linear(self.hidden_size, num_heads * head_dim, bias=False) for _ in range(self.num_memories)
        ])

        self.router = nn.Linear(self.hidden_size, self.num_memories, bias=False)
        if getattr(config, 'stable_moe_stage', 1) == 2:
            for p in self.router.parameters():
                p.requires_grad = False

        if self.shared_memory:
            self.shared_shortconv_k = nn.Conv1d(self.hidden_size, self.hidden_size, conv_kernel, padding=pad, groups=self.hidden_size, bias=False)
            self.shared_shortconv_v = nn.Conv1d(self.hidden_size, self.hidden_size, conv_kernel, padding=pad, groups=self.hidden_size, bias=False)
            self.shared_k_proj = nn.Linear(self.hidden_size, num_heads * head_dim, bias=False)
            self.shared_v_proj = nn.Linear(self.hidden_size, num_heads * head_dim, bias=False)
            self.shared_alpha_down = nn.Linear(self.hidden_size, head_dim, bias=False)
            self.shared_alpha_up = nn.Linear(head_dim, num_heads * head_dim, bias=False)
            self.shared_beta_proj = nn.Linear(self.hidden_size, num_heads, bias=False)
            self.shared_gate_proj = nn.Linear(self.hidden_size, num_heads * head_dim, bias=False)

    def _apply_shortconv(self, x, conv):
        x_t = x.transpose(1, 2)
        x_t = self.act_fn(x_t)
        x_t = conv(x_t)
        return x_t.transpose(1, 2)

    def _recurrent_step(self, q, k, v, alpha, beta, state):
        batch, nh, hd = q.shape
        k_u = k.unsqueeze(-1)
        v_u = v.unsqueeze(-1)
        q_u = q.unsqueeze(-1)
        alpha_diag = alpha.unsqueeze(-1)
        state_scaled = state * alpha_diag
        k_state = torch.matmul(k_u.transpose(-2, -1), state_scaled)
        beta_k = beta.view(batch, nh, 1, 1)
        state = state_scaled - beta_k * (k_u @ k_state)
        state = state + beta_k * (k_u @ v_u.transpose(-2, -1))
        out = (state @ q_u).squeeze(-1)
        return out, state

    def _run_memory(self, hidden_states, shortconv_k, shortconv_v, k_proj, v_proj,
                    alpha_down, alpha_up, beta_proj, gate_proj, q, state):
        bsz, sl, _ = hidden_states.shape
        if state is None:
            state = torch.zeros(bsz, self.num_heads, self.head_dim, self.head_dim,
                              device=hidden_states.device, dtype=hidden_states.dtype)
        k_sconv = self._apply_shortconv(hidden_states, shortconv_k)
        v_sconv = self._apply_shortconv(hidden_states, shortconv_v)
        k = k_proj(k_sconv).view(bsz, sl, self.num_heads, self.head_dim)
        v = v_proj(v_sconv).view(bsz, sl, self.num_heads, self.head_dim)
        k = F.normalize(k, dim=-1)

        h = alpha_down(hidden_states)
        h = self.act_fn(h)
        h = alpha_up(h)
        alpha = F.sigmoid(h).view(bsz, sl, self.num_heads, self.head_dim)
        beta = F.sigmoid(beta_proj(hidden_states))

        if sl == 1 and state is not None:
            out_i, new_state = self._recurrent_step(q[:, 0], k[:, 0], v[:, 0],
                                                      alpha[:, 0], beta[:, 0], state)
            out_i = self.output_norm(out_i)
            gate = F.sigmoid(gate_proj(hidden_states[:, 0])).view(bsz, self.num_heads, self.head_dim)
            out_i = out_i * gate
            return out_i.reshape(bsz, 1, -1), new_state

        outputs = []
        cur_state = state
        for t in range(sl):
            o_t, cur_state = self._recurrent_step(q[:, t], k[:, t], v[:, t],
                                                    alpha[:, t], beta[:, t], cur_state)
            outputs.append(o_t)
        out = torch.stack(outputs, dim=1)
        out = self.output_norm(out)
        gate = F.sigmoid(gate_proj(hidden_states)).view(bsz, sl, self.num_heads, self.head_dim)
        out = out * gate
        return out.reshape(bsz, sl, -1), cur_state

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask=None,
        position_ids=None,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        output_attentions: bool = False,
        use_cache: bool = False,
    ):
        bsz, sl, h = hidden_states.shape
        q_sconv = self._apply_shortconv(hidden_states, self.shortconv_q)
        q = self.q_proj(q_sconv).view(bsz, sl, self.num_heads, self.head_dim)
        q = F.normalize(q, dim=-1)

        past_states = None
        if past_key_value is not None and past_key_value[0] is not None:
            past_states = past_key_value[0]

        router_logits = self.router(hidden_states)
        routing_weights = F.softmax(router_logits, dim=-1, dtype=torch.float)
        _, topk_idx = torch.topk(routing_weights, self.active_memories, dim=-1)
        topk_mask = torch.zeros(bsz, sl, self.num_memories, device=hidden_states.device)
        topk_mask.scatter_(-1, topk_idx, 1.0)

        mem_outputs = []
        new_states = []
        for m in range(self.num_memories):
            mem_mask = topk_mask[:, :, m]
            mem_state = past_states[m] if past_states is not None and m < past_states.shape[0] else None

            out_m, new_state_m = self._run_memory(
                hidden_states, self.shortconv_k[m], self.shortconv_v[m],
                self.k_proj[m], self.v_proj[m],
                self.alpha_down[m], self.alpha_up[m],
                self.beta_proj[m], self.gate_proj[m], q, mem_state
            )
            mem_outputs.append(out_m * mem_mask.unsqueeze(-1))
            new_states.append(new_state_m)

        if self.shared_memory:
            shared_state = past_states[self.num_memories] if past_states is not None and past_states.shape[0] > self.num_memories else None
            out_s, new_state_s = self._run_memory(
                hidden_states, self.shared_shortconv_k, self.shared_shortconv_v,
                self.shared_k_proj, self.shared_v_proj,
                self.shared_alpha_down, self.shared_alpha_up,
                self.shared_beta_proj, self.shared_gate_proj, q, shared_state
            )
            mem_outputs.append(out_s)
            new_states.append(new_state_s)

        total_out = sum(mem_outputs)
        total_out = self.o_proj(total_out)

        all_states = torch.stack(new_states, dim=0)
        present = (all_states,) if use_cache else None
        return total_out, None, present


class QuantileBalancingRouter(nn.Module):
    """MoE router with Quantile Balancing for load balance.

    Alternating quantile algorithm (J. Su) computes per-expert biases
    that equalize token assignment — hyperparameter-free.
    """

    def __init__(self, config, num_experts):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = config.num_experts_per_tok
        self.qb_iterations = config.qb_iterations
        self.quantile_balancing = config.quantile_balancing

        self.router = nn.Linear(config.hidden_size, num_experts, bias=False)

        if hasattr(config, 'stable_moe_stage') and config.stable_moe_stage == 2:
            for p in self.router.parameters():
                p.requires_grad = False

    def _quantile_bias(self, scores):
        m, n = scores.shape
        k = self.top_k
        beta = torch.zeros(1, n, device=scores.device, dtype=scores.dtype)
        q = 1.0 - k / n
        for _ in range(self.qb_iterations):
            alpha = torch.quantile(scores - beta, q, dim=1, keepdim=True)
            beta = torch.quantile(scores - alpha, q, dim=0, keepdim=True)
        return beta

    def forward(self, x):
        router_logits = self.router(x)
        if self.quantile_balancing and self.training:
            bias = self._quantile_bias(router_logits)
            biased_logits = router_logits - bias
        else:
            biased_logits = router_logits
        routing_weights = F.softmax(biased_logits, dim=-1, dtype=torch.float)
        routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)
        routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True)
        return routing_weights.to(x.dtype), selected_experts, router_logits


class ScrapeGoatTrackAMoE(nn.Module):
    """Track A MoE: fused 3D expert weights with shared expert, 512 experts, QB routing."""

    def __init__(self, config):
        super().__init__()
        self.num_experts = config.track_a_num_experts
        self.moe_intermediate_size = config.track_a_moe_intermediate_size
        self.hidden_size = config.hidden_size
        self.top_k = config.num_experts_per_tok
        self.act_fn = lambda x: activation_fn(x, config.hidden_act)

        self.router = QuantileBalancingRouter(config, self.num_experts)
        self.gate_up_weight = nn.Parameter(torch.empty(self.num_experts, 2 * self.moe_intermediate_size, self.hidden_size))
        self.down_weight = nn.Parameter(torch.empty(self.num_experts, self.hidden_size, self.moe_intermediate_size))

        self.shared_gate = nn.Linear(self.hidden_size, self.moe_intermediate_size, bias=False)
        self.shared_up = nn.Linear(self.hidden_size, self.moe_intermediate_size, bias=False)
        self.shared_down = nn.Linear(self.moe_intermediate_size, self.hidden_size, bias=False)
        self.shared_expert_gate = nn.Linear(self.hidden_size, 1, bias=False)

    def forward(self, x):
        bsz, seq_len, h = x.shape
        x_flat = x.view(-1, h)

        routing_weights, selected_experts, router_logits = self.router(x_flat)

        final = torch.zeros_like(x_flat)

        for k in range(self.top_k):
            e_idx = selected_experts[:, k]
            w = routing_weights[:, k:k + 1]

            fused = self.gate_up_weight[e_idx]
            gate_w = fused[:, :self.moe_intermediate_size]
            up_w = fused[:, self.moe_intermediate_size:]

            gate_out = torch.bmm(gate_w, x_flat.unsqueeze(-1)).squeeze(-1)
            up_out = torch.bmm(up_w, x_flat.unsqueeze(-1)).squeeze(-1)
            expert_h = self.act_fn(gate_out) * up_out

            down_w = self.down_weight[e_idx]
            expert_out = torch.bmm(down_w, expert_h.unsqueeze(-1)).squeeze(-1)
            final += w * expert_out

        shared_gate_val = torch.sigmoid(self.shared_expert_gate(x_flat))
        shared_h = self.act_fn(self.shared_gate(x_flat)) * self.shared_up(x_flat)
        shared_out = self.shared_down(shared_h)
        final = final + shared_gate_val * shared_out

        return final.view(bsz, seq_len, h), router_logits


class ScrapeGoatTrackBMoE(nn.Module):
    """Track B MoE: individual experts stacked into 3D, 192 experts + shared MLP, QB routing."""

    def __init__(self, config):
        super().__init__()
        self.num_experts = config.track_b_num_experts
        self.moe_intermediate_size = config.track_b_moe_intermediate_size
        self.hidden_size = config.hidden_size
        self.top_k = config.num_experts_per_tok
        self.act_fn = lambda x: activation_fn(x, config.hidden_act)
        self.router = QuantileBalancingRouter(config, self.num_experts)

        self.gate_proj_weight = nn.Parameter(torch.empty(self.num_experts, self.moe_intermediate_size, self.hidden_size))
        self.up_proj_weight = nn.Parameter(torch.empty(self.num_experts, self.moe_intermediate_size, self.hidden_size))
        self.down_proj_weight = nn.Parameter(torch.empty(self.num_experts, self.hidden_size, self.moe_intermediate_size))

        self.shared_gate = nn.Linear(self.hidden_size, self.moe_intermediate_size, bias=False)
        self.shared_up = nn.Linear(self.hidden_size, self.moe_intermediate_size, bias=False)
        self.shared_down = nn.Linear(self.moe_intermediate_size, self.hidden_size, bias=False)

    def forward(self, x):
        bsz, seq_len, h = x.shape
        x_flat = x.view(-1, h)

        routing_weights, selected_experts, router_logits = self.router(x_flat)

        final = torch.zeros_like(x_flat)

        for k in range(self.top_k):
            e_idx = selected_experts[:, k]
            w = routing_weights[:, k:k + 1]

            gate_w = self.gate_proj_weight[e_idx]
            up_w = self.up_proj_weight[e_idx]
            gate_out = torch.bmm(gate_w, x_flat.unsqueeze(-1)).squeeze(-1)
            up_out = torch.bmm(up_w, x_flat.unsqueeze(-1)).squeeze(-1)
            expert_h = self.act_fn(gate_out) * up_out

            down_w = self.down_proj_weight[e_idx]
            expert_out = torch.bmm(down_w, expert_h.unsqueeze(-1)).squeeze(-1)
            final += w * expert_out

        shared_h = self.act_fn(self.shared_gate(x_flat)) * self.shared_up(x_flat)
        shared_out = self.shared_down(shared_h)
        final = final + shared_out

        return final.view(bsz, seq_len, h), router_logits


class ScrapeGoatTrackBDenseFFN(nn.Module):
    """Track B dense FFN for layer 0 (from intermediate_size)."""

    def __init__(self, config):
        super().__init__()
        self.act_fn = lambda x: activation_fn(x, config.hidden_act)
        inter = config.track_b_intermediate_size
        self.gate_proj = nn.Linear(config.hidden_size, inter, bias=False)
        self.up_proj = nn.Linear(config.hidden_size, inter, bias=False)
        self.down_proj = nn.Linear(inter, config.hidden_size, bias=False)

    def forward(self, x):
        return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))


class BlockAttnRes(nn.Module):
    """Block Attention Residuals — softmax attention over depth.

    Partitions layers into N blocks. Each layer attends over prior
    block representations via a learned pseudo-query.
    """

    def __init__(self, config, layer_idx):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.layer_idx = layer_idx
        self.num_blocks = config.attn_res_blocks
        self.block_size = max(1, config.num_hidden_layers // self.num_blocks)
        self.block_idx = layer_idx // self.block_size

        self.pseudo_query = nn.Parameter(torch.zeros(self.hidden_size))
        self.query_norm = ScrapeGoatRMSNorm(1)

        self.block_norm = ScrapeGoatRMSNorm(self.hidden_size)

    def forward(self, hidden_states, block_reps):
        if block_reps is None or len(block_reps) == 0:
            return hidden_states

        q = self.pseudo_query.unsqueeze(0).unsqueeze(0)
        stacked = torch.stack(block_reps, dim=1)
        bsz, n_blocks, _ = stacked.shape

        keys = self.block_norm(stacked)
        query = self.query_norm(q)

        attn_weights = torch.matmul(query, keys.transpose(-2, -1)) / math.sqrt(self.hidden_size)
        attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(stacked.dtype)

        depth_out = torch.matmul(attn_weights, stacked).squeeze(1)
        return hidden_states + depth_out


class ScrapeGoatParallelDecoderLayer(nn.Module):
    """Decoder layer with parallel dual-track attention/MoE and v2 enhancements.

    v2 additions:
    - Track A: 3:1 KDA:GQA interleaving
    - Track B: always GQA
    - Quantile Balancing for MoE routing
    - Block Attention Residuals
    """

    def __init__(self, config: ScrapeGoatConfig, layer_idx: int):
        super().__init__()
        self.layer_idx = layer_idx
        self.hidden_size = config.hidden_size

        self.input_layernorm = ScrapeGoatRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.post_attention_layernorm = ScrapeGoatRMSNorm(config.hidden_size, eps=config.rms_norm_eps)

        use_kda = layer_idx not in config.kda_gqa_layers
        if use_kda:
            if getattr(config, 'mom_enabled', False):
                self.track_a_attn = ScrapeGoatMoMKDA(
                    config,
                    num_heads=config.track_a_num_attention_heads,
                    head_dim=config.kda_head_dim,
                )
            else:
                self.track_a_attn = ScrapeGoatKDA(
                    config,
                    num_heads=config.track_a_num_attention_heads,
                    head_dim=config.kda_head_dim,
                )
        else:
            self.track_a_attn = ScrapeGoatSourceAttention(
                config,
                num_heads=config.track_a_num_attention_heads,
                num_kv_heads=config.track_a_num_key_value_heads,
                head_dim=config.track_a_head_dim,
            )
        self.track_b_attn = ScrapeGoatSourceAttention(
            config,
            num_heads=config.track_b_num_attention_heads,
            num_kv_heads=config.track_b_num_key_value_heads,
            head_dim=config.track_b_head_dim,
        )
        self.attn_track_gate = nn.Linear(config.hidden_size, 2, bias=False)

        self.track_a_moe = ScrapeGoatTrackAMoE(config)
        if layer_idx == 0:
            self.track_b_dense_ffn = ScrapeGoatTrackBDenseFFN(config)
            self.track_b_moe = None
        else:
            self.track_b_moe = ScrapeGoatTrackBMoE(config)
        self.moe_track_gate = nn.Linear(config.hidden_size, 2, bias=False)

        if config.attn_residual:
            self.attn_res = BlockAttnRes(config, layer_idx)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        output_attentions: bool = False,
        use_cache: bool = False,
        block_reps: Optional[List[torch.Tensor]] = None,
    ) -> Tuple:
        residual = hidden_states
        hidden_states = self.input_layernorm(hidden_states)

        track_a_past = past_key_value[0] if past_key_value is not None else None
        track_b_past = past_key_value[1] if past_key_value is not None else None

        track_a_attn_out, track_a_attn_w, track_a_pk = self.track_a_attn(
            hidden_states, attention_mask, position_ids,
            track_a_past, output_attentions, use_cache,
        )
        track_b_attn_out, track_b_attn_w, track_b_pk = self.track_b_attn(
            hidden_states, attention_mask, position_ids,
            track_b_past, output_attentions, use_cache,
        )

        attn_gate = torch.sigmoid(self.attn_track_gate(hidden_states))
        hidden_states = residual + attn_gate[:, :, :1] * track_a_attn_out + attn_gate[:, :, 1:] * track_b_attn_out

        if hasattr(self, 'attn_res'):
            hidden_states = self.attn_res(hidden_states, block_reps)

        residual = hidden_states
        hidden_states = self.post_attention_layernorm(hidden_states)

        track_a_moe_out, track_a_router = self.track_a_moe(hidden_states)
        if self.track_b_moe is not None:
            track_b_moe_out, track_b_router = self.track_b_moe(hidden_states)
        else:
            track_b_moe_out = self.track_b_dense_ffn(hidden_states)

        moe_gate = torch.sigmoid(self.moe_track_gate(hidden_states))
        hidden_states = residual + moe_gate[:, :, :1] * track_a_moe_out + moe_gate[:, :, 1:] * track_b_moe_out

        outputs = (hidden_states,)
        if output_attentions:
            outputs += (track_a_attn_w, track_b_attn_w)
        if use_cache:
            outputs += (track_a_pk, track_b_pk)

        return outputs


class ScrapeGoatPreTrainedModel(PreTrainedModel):
    config_class = ScrapeGoatConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["ScrapeGoatParallelDecoderLayer"]
    _skip_keys_device_placement = "past_key_values"
    _supports_flash_attn_2 = True
    _supports_sdpa = True
    _supports_cache_class = True

    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)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.Parameter):
            module.data.normal_(mean=0.0, std=std)


class ScrapeGoatModel(ScrapeGoatPreTrainedModel):
    def __init__(self, config: ScrapeGoatConfig):
        super().__init__(config)
        self.padding_idx = config.pad_token_id
        self.vocab_size = config.vocab_size

        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
        self.layers = nn.ModuleList(
            [ScrapeGoatParallelDecoderLayer(config, idx) for idx in range(config.num_hidden_layers)]
        )
        self.norm = ScrapeGoatRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
        self.gradient_checkpointing = False
        self.post_init()

    def get_input_embeddings(self):
        return self.embed_tokens

    def set_input_embeddings(self, value):
        self.embed_tokens = value

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        use_cache: bool = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, BaseModelOutputWithPast]:
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        use_cache = use_cache if use_cache is not None else self.config.use_cache
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            batch_size, seq_length = input_ids.shape
        elif inputs_embeds is not None:
            batch_size, seq_length = inputs_embeds.shape[:2]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        if self.gradient_checkpointing and self.training:
            if use_cache:
                logger.warning_once(
                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."
                )
                use_cache = False

        past_key_values_length = 0
        if past_key_values is not None:
            past_key_values_length = past_key_values[0][1][0].shape[-2]

        if position_ids is None:
            device = input_ids.device if input_ids is not None else inputs_embeds.device
            position_ids = torch.arange(
                past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device
            )
            position_ids = position_ids.unsqueeze(0)

        if inputs_embeds is None:
            inputs_embeds = self.embed_tokens(input_ids)

        if attention_mask is None:
            attention_mask = torch.ones((batch_size, seq_length + past_key_values_length), dtype=torch.bool, device=inputs_embeds.device)
        attention_mask = self._prepare_decoder_attention_mask(
            attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length
        )

        hidden_states = inputs_embeds

        all_hidden_states = () if output_hidden_states else None
        all_self_attns = () if output_attentions else None
        next_decoder_cache = () if use_cache else None

        block_reps = []
        block_size = max(1, self.config.num_hidden_layers // self.config.attn_res_blocks) if self.config.attn_residual else None
        prev_block_idx = -1

        for layer_idx, decoder_layer in enumerate(self.layers):
            if output_hidden_states:
                all_hidden_states += (hidden_states,)

            layer_past = past_key_values[layer_idx] if past_key_values is not None else None

            layer_outputs = decoder_layer(
                hidden_states,
                attention_mask=attention_mask,
                position_ids=position_ids,
                past_key_value=layer_past,
                output_attentions=output_attentions,
                use_cache=use_cache,
                block_reps=block_reps if self.config.attn_residual else None,
            )

            hidden_states = layer_outputs[0]

            if use_cache:
                next_decoder_cache += ((layer_outputs[-2], layer_outputs[-1]),)

            if output_attentions:
                all_self_attns += (layer_outputs[1], layer_outputs[2])

            if block_size is not None:
                current_block = decoder_layer.layer_idx // block_size
                if current_block > prev_block_idx:
                    block_reps.append(hidden_states.detach().mean(dim=1, keepdim=False))
                    prev_block_idx = current_block

        hidden_states = self.norm(hidden_states)

        if output_hidden_states:
            all_hidden_states += (hidden_states,)

        next_cache = next_decoder_cache if use_cache else None
        if not return_dict:
            return tuple(
                v
                for v in [hidden_states, next_cache, all_hidden_states, all_self_attns]
                if v is not None
            )
        return BaseModelOutputWithPast(
            last_hidden_state=hidden_states,
            past_key_values=next_cache,
            hidden_states=all_hidden_states,
            attentions=all_self_attns,
        )

    def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):
        combined_attention_mask = None
        if input_shape[-1] > 0:
            combined_attention_mask = _make_causal_mask(
                input_shape, inputs_embeds.dtype, past_key_values_length=past_key_values_length
            ).to(inputs_embeds.device)

        if attention_mask is not None:
            padded_attention_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(
                inputs_embeds.device
            )
            combined_attention_mask = (
                padded_attention_mask if combined_attention_mask is None else combined_attention_mask + padded_attention_mask
            )

        return combined_attention_mask


class ScrapeGoatForCausalLM(ScrapeGoatPreTrainedModel):
    _tied_weights_keys = ["lm_head.weight"]

    def __init__(self, config: ScrapeGoatConfig):
        super().__init__(config)
        self.model = ScrapeGoatModel(config)
        self.vocab_size = config.vocab_size
        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)

        self.dspark_attention = DSparkAttention(config)
        self.dspark_markov_head = DSparkMarkovHead(config)
        self.dspark_confidence_head = DSparkConfidenceHead(config)

        self.dspark_target_layer_ids = set(config.dspark_target_layer_ids) if config.dspark_target_layer_ids else set()

        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, value):
        self.lm_head = value

    def set_decoder(self, decoder):
        self.model = decoder

    def get_decoder(self):
        return self.model

    def forward(
        self,
        input_ids: torch.LongTensor = None,
        attention_mask: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: bool = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        router_logits: Optional[bool] = None,
    ) -> Union[Tuple, CausalLMOutputWithPast]:
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        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=output_hidden_states,
            return_dict=return_dict,
        )

        hidden_states = outputs[0]
        logits = self.lm_head(hidden_states)
        logits = logits.float()

        loss = None
        if labels is not None:
            shift_logits = logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            loss_fct = nn.CrossEntropyLoss()
            shift_logits = shift_logits.view(-1, self.vocab_size)
            shift_labels = shift_labels.view(-1)
            shift_labels = shift_labels.to(shift_logits.device)
            loss = loss_fct(shift_logits, shift_labels)

        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,
        )

    def prepare_inputs_for_generation(
        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs
    ):
        if past_key_values is not None:
            input_ids = input_ids[:, -1:]
        position_ids = kwargs.get("position_ids", None)
        if attention_mask is not None and position_ids is None:
            position_ids = attention_mask.long().cumsum(-1) - 1
            position_ids.masked_fill_(attention_mask == 0, 1)
            if past_key_values:
                position_ids = position_ids[:, -1].unsqueeze(-1)

        if inputs_embeds is not None and past_key_values is None:
            model_inputs = {"inputs_embeds": inputs_embeds}
        else:
            model_inputs = {"input_ids": input_ids}
        model_inputs.update(
            {
                "position_ids": position_ids,
                "past_key_values": past_key_values,
                "use_ebd": kwargs.get("use_ebd"),
                "attention_mask": attention_mask,
            }
        )
        return model_inputs

    @staticmethod
    def _reorder_cache(past_key_values, beam_idx):
        reordered_past = ()
        for past in past_key_values:
            reordered_past += (
                past[0].index_select(0, beam_idx),
                past[1].index_select(0, beam_idx),
            )
        return reordered_past


def _make_causal_mask(input_ids_shape, dtype, past_key_values_length=0):
    bsz, tgt_len = input_ids_shape
    mask = torch.full((tgt_len, tgt_len), torch.tensor(torch.finfo(dtype).min), device="cpu")
    mask_cond = torch.arange(mask.size(-1))
    mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)
    mask = mask.to(dtype)
    if past_key_values_length > 0:
        mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype), mask], dim=-1)
    return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)


def _expand_mask(mask, dtype, tgt_len=None):
    bsz, src_len = mask.shape
    tgt_len = tgt_len if tgt_len is not None else src_len
    expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
    inverted_mask = 1.0 - expanded_mask
    return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)