dspark_components.py

"""
DSpark-specific components for ScrapeGoat model.
Implements DSpark attention, Markov heads, and related components.
"""

import math
import torch
import torch.nn as nn
from typing import Optional, Tuple


def get_dspark_topk_idxs(window_size: int, bsz: int, block_size: int, start_pos: int):
    """
    Get top-k indices for DSpark attention.
    Based on the implementation from DeepSeek-V4-Pro-DSpark.
    """
    assert start_pos > 0
    # Create tensor: [0, 1, ..., min(window_size, start_pos+1)-1, window_size, window_size+1, ..., window_size+block_size-1]
    idx1 = torch.arange(min(window_size, start_pos + 1), device='cpu')  # Will be moved to correct device later
    idx2 = window_size + torch.arange(block_size, device='cpu')
    matrix = torch.cat([idx1, idx2])
    return matrix.int().view(1, 1, -1).expand(bsz, block_size, -1)


class DSparkAttention(nn.Module):
    """
    DSpark Attention mechanism from DeepSeek-V4-Pro-DSpark.
    Combines window-based attention with compressed attention for efficient long-context processing.
    """

    def __init__(self, config):
        super().__init__()
        self.config = config
        self.hidden_size = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.head_dim = self.hidden_size // self.num_heads
        self.num_key_value_heads = config.num_key_value_heads
        self.num_key_value_groups = self.num_heads // self.num_key_value_heads
        self.max_position_embeddings = config.max_position_embeddings
        self.rope_theta = getattr(config, 'rope_theta', 10000.0)
        self.is_causal = True
        self.attn_sink = nn.Parameter(torch.zeros(self.num_heads))  # Learnable attention sink bias

        # Projections
        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)
        self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
        self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)
        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias)

        # For RoPE
        self.rotary_emb = None  # Will be set by the model or created internally if needed

        # DSpark-specific parameters
        self.window_size = getattr(config, 'window_size', 128)
        self.compress_ratio = getattr(config, 'compress_ratio', 4)
        self.has_indexer = self.compress_ratio == 4

        # KV cache will be managed externally
        self.kv_cache = None

    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,
        main_x: Optional[torch.Tensor] = None,  # This is the key DSpark input - features from target layers
        output_attentions: bool = False,
        use_cache: bool = False,
        **kwargs,
    ) -> tuple:
        """
        DSpark Attention forward pass.

        Args:
            hidden_states: Input tensor [batch_size, seq_len, hidden_size]
            attention_mask: Attention mask
            position_ids: Position IDs
            past_key_value: Cached key/value states
            main_x: Features from target layers (specific to DSpark)
            output_attentions: Whether to return attention weights
            use_cache: Whether to use caching

        Returns:
            tuple of (output, attention_weights, present_key_value)
        """
        bsz, q_len, _ = hidden_states.size()

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

        # Reshape for multi-head attention
        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)

        # Handle past key value
        kv_seq_len = key_states.shape[-2]
        if past_key_value is not None:
            kv_seq_len += past_key_value[0].shape[-2]
            key_states = torch.cat([past_key_value[0], key_states], dim=2)
            value_states = torch.cat([past_key_value[1], value_states], dim=2)

        # Apply RoPE if available (simplified - in practice would use precomputed freqs)
        # For now, we'll skip RoPE implementation details and focus on DSpark logic

        # Repeat k/v heads if n_kv_heads < n_heads
        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)

        # DSpark-specific logic: if we have main_x, use DSpark attention
        if main_x is not None:
            attn_output = self._dspark_attention_forward(
                query_states, key_states, value_states,
                attention_mask, position_ids, past_key_value, main_x
            )
        else:
            # Standard attention
            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 = nn.functional.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().view(bsz, q_len, self.hidden_size)

        attn_output = self.o_proj(attn_output)

        if not output_attentions:
            attn_weights = None

        present_key_value = (key_states, value_states) if use_cache else None

        return attn_output, attn_weights, present_key_value

    def _dspark_attention_forward(
        self,
        query_states: torch.Tensor,
        key_states: torch.Tensor,
        value_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor],
        position_ids: Optional[torch.LongTensor],
        past_key_value: Optional[Tuple[torch.Tensor]],
        main_x: torch.Tensor
    ) -> torch.Tensor:
        """
        DSpark-specific attention computation that combines:
        1. Window attention on recent tokens from main_x
        2. Compressed attention on older tokens from main_x
        3. Standard attention on current hidden_states
        """
        bsz, q_len, _ = query_states.shape[:2]  # [b, q_len, h*d] -> reshape to [b, q_len, h, d]

        # Reshape query states for processing
        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim)

        # Process main_x to get key/value states for DSpark
        # Main x comes from target layers, so we need to project it to K/V space
        main_key_states = self.k_proj(main_x)
        main_value_states = self.v_proj(main_x)

        # Reshape main key/value states
        main_key_states = main_key_states.view(main_x.size(0), main_x.size(1),
                                              self.num_key_value_heads, self.head_dim).transpose(1, 2)
        main_value_states = main_value_states.view(main_x.size(0), main_x.size(1),
                                                  self.num_key_value_heads, self.head_dim).transpose(1, 2)

        # Repeat k/v heads for main_x
        main_key_states = main_key_states.repeat_interleave(self.num_key_value_groups, dim=1)
        main_value_states = main_value_states.repeat_interleave(self.num_key_value_groups, dim=1)

        # Current sequence length from main_x
        main_seq_len = main_x.size(1)

        # Calculate effective lengths for window and compression
        effective_window_size = min(self.window_size, main_seq_len)

        # Window attention: attend to recent tokens in main_x
        if effective_window_size > 0:
            window_start = max(0, main_seq_len - effective_window_size)
            window_keys = main_key_states[:, :, window_start:main_seq_len, :]
            window_values = main_value_states[:, :, window_start:main_seq_len, :]

            # Compute window attention scores
            window_q = query_states  # [b, q_len, h, d]
            window_k = window_keys   # [b, h, kv_len, d]

            window_scores = torch.matmul(window_q, window_k.transpose(-2, -1)) / math.sqrt(self.head_dim)

            # Apply causal mask if needed
            if self.is_causal:
                q_positions = torch.arange(q_len, device=query_states.device).unsqueeze(1)
                k_positions = torch.arange(window_start, main_seq_len, device=query_states.device).unsqueeze(0)
                causal_mask = q_positions >= (k_positions - window_start)
                causal_mask = causal_mask.unsqueeze(0).unsqueeze(1)  # [1, 1, q_len, kv_len]
                window_scores = window_scores.masked_fill(~causal_mask, float('-inf'))

            # Apply attention sink bias
            window_scores = window_scores + self.attn_sink.view(1, -1, 1, 1)

            window_attn_weights = nn.functional.softmax(window_scores, dim=-1, dtype=torch.float32)
            window_attn_output = torch.matmul(window_attn_weights.to(window_values.dtype), window_values)
        else:
            window_attn_output = torch.zeros_like(query_states)

        # Compressed attention: attend to compressed representation of older tokens
        # This is a simplified version - full implementation would use the compressor/indexer
        if main_seq_len > self.window_size:
            # For simplicity, we'll just use average pooling as compression
            # In practice, this would use the Compressor and Indexer modules
            remaining_len = max(0, main_seq_len - self.window_size)
            if remaining_len > 0:
                # Simple average compression (placeholder)
                compressed_k = torch.mean(main_key_states[:, :, :main_seq_len - self.window_size, :], dim=2, keepdim=True)
                compressed_v = torch.mean(main_value_states[:, :, :main_seq_len - self.window_size, :], dim=2, keepdim=True)

                # Expand to match heads
                compressed_k = compressed_k.expand(-1, self.num_heads, -1, -1)
                compressed_v = compressed_v.expand(-1, self.num_heads, -1, -1)

                # Compute compressed attention
                compressed_q = query_states
                compressed_scores = torch.matmul(compressed_q, compressed_k.transpose(-2, -1)) / math.sqrt(self.head_dim)
                compressed_attn_weights = nn.functional.softmax(compressed_scores, dim=-1, dtype=torch.float32)
                compressed_attn_output = torch.matmul(compressed_attn_weights.to(compressed_v.dtype), compressed_v)
            else:
                compressed_attn_output = torch.zeros_like(query_states)
        else:
            compressed_attn_output = torch.zeros_like(query_states)

        # Standard attention on current hidden states (if needed)
        # In DSpark, this might be skipped or weighted differently
        std_attn_output = torch.zeros_like(query_states)  # Placeholder

        # Combine outputs (in practice, this would be learned weights)
        # For now, simple sum
        combined_output = window_attn_output + compressed_attn_output + std_attn_output

        # Reshape back
        combined_output = combined_output.transpose(1, 2).contiguous().view(bsz, q_len, self.hidden_size)

        return combined_output


class DSparkMarkovHead(nn.Module):
    """
    DSpark Markov Head for next token prediction based on Markov chains.
    """

    def __init__(self, config):
        super().__init__()
        self.vocab_size = config.vocab_size
        self.markov_rank = getattr(config, 'dspark_markov_rank', 256)

        # Markov transition matrices
        self.markov_w1 = nn.Embedding(self.vocab_size, self.markov_rank)
        self.markov_w2 = nn.Linear(self.markov_rank, self.vocab_size, bias=False)

    def forward(self, token_ids: torch.Tensor) -> torch.Tensor:
        """
        Compute Markov-based next token logits.

        Args:
            token_ids: Input token IDs [batch_size, seq_len]

        Returns:
            logits: Next token logits [batch_size, seq_len, vocab_size]
        """
        # Embed tokens
        embed = self.markov_w1(token_ids)  # [batch_size, seq_len, markov_rank]

        # Project to vocabulary space
        logits = self.markov_w2(embed)  # [batch_size, seq_len, vocab_size]

        return logits


class DSparkConfidenceHead(nn.Module):
    """
    DSpark Confidence Head for scoring prediction confidence.
    """

    def __init__(self, config):
        super().__init__()
        hidden_size = getattr(config, 'hidden_size', 4096)
        markov_rank = getattr(config, 'dspark_markov_rank', 256)
        input_dim = hidden_size + markov_rank

        self.proj = nn.Linear(input_dim, 1, bias=False)

    def forward(self, hidden: torch.Tensor, markov_embed: torch.Tensor) -> torch.Tensor:
        """
        Compute confidence score.

        Args:
            hidden: Hidden states from model [batch_size, seq_len, hidden_size]
            markov_embed: Markov embeddings [batch_size, seq_len, markov_rank]

        Returns:
            confidence: Confidence scores [batch_size, seq_len]
        """
        # Concatenate hidden states and Markov embeddings
        combined = torch.cat([hidden, markov_embed], dim=-1)  # [batch_size, seq_len, hidden_size + markov_rank]

        # Project to single dimension
        confidence = self.proj(combined).squeeze(-1)  # [batch_size, seq_len]

        return confidence


def sample(logits, temperature: float = 1.0):
    """
    Sample from logits using Gumbel-max trick.
    """
    if temperature == 0:
        return logits.argmax(dim=-1)
    logits = logits / max(temperature, 1e-5)
    probs = torch.softmax(logits, dim=-1, dtype=torch.float32)
    return probs.div_(torch.empty_like(probs).exponential_(1)).argmax(dim=-1)