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