""" Context Window — Janela deslizante com cache KV para sequências longas. Implementa três estratégias (ver docs/MATH_ANALYSIS.md seção 3.4): - Sliding Window Pura - Sink + Sliding (StreamingLLM) — RECOMENDADO - Attention Recomputation (sem cache, máxima precisão) Para CNN-BiGRU cooperativo, o contexto é gerenciado em modo: - Encoder (treino/classificação): processa janela inteira, BiGRU bidirecional - Decoder (geração): cache KV para CausalSelfAttention + GRU unidirecional Autor: CNN-BiGRU Project """ from __future__ import annotations import logging from dataclasses import dataclass, field from typing import Dict, List, Optional, Tuple import torch import torch.nn as nn logger = logging.getLogger(__name__) # ============================================================================ # Configuração # ============================================================================ @dataclass class ContextWindowConfig: """Configuração da janela de contexto.""" # Comprimento máximo da janela (em tokens) max_window: int = 512 # Estratégia de evicção: "sliding" | "sink_sliding" | "recompute" eviction_strategy: str = "sink_sliding" # Número de tokens "sink" (apenas para sink_sliding) — tipicamente BOS + system prompt n_sink_tokens: int = 4 # Dimensão do embedding (para alocar cache KV) embed_dim: int = 256 # Número de cabeças (para cache KV em multi-head attention) n_heads: int = 4 # Dimensão por cabeça head_dim: Optional[int] = None # default: embed_dim // n_heads # Número de camadas (para cache KV multicamada) n_layers: int = 2 # Device padrão device: str = "cpu" # Dtype do cache (None = manter fp32) dtype: Optional[torch.dtype] = None # ============================================================================ # KV Cache # ============================================================================ class KVCache: """ Cache de pares (Key, Value) para CausalSelfAttention. Shape por camada: K: [n_heads, seq_cached, head_dim] V: [n_heads, seq_cached, head_dim] Em modo batch: K: [batch, n_heads, seq_cached, head_dim] V: [batch, n_heads, seq_cached, head_dim] """ def __init__(self, n_layers: int, batch_size: int, n_heads: int, head_dim: int, device: torch.device, dtype: torch.dtype): self.n_layers = n_layers self.batch_size = batch_size self.n_heads = n_heads self.head_dim = head_dim self.device = device self.dtype = dtype # Pre-alocar listas vazias; preencher lazy no primeiro update self.keys: List[Optional[torch.Tensor]] = [None] * n_layers self.values: List[Optional[torch.Tensor]] = [None] * n_layers def update(self, layer_idx: int, new_k: torch.Tensor, new_v: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """ Atualiza o cache da camada layer_idx com novos K, V. Args: layer_idx: índice da camada [0, n_layers) new_k: [batch, n_heads, seq_new, head_dim] new_v: [batch, n_heads, seq_new, head_dim] Returns: (cached_k, cached_v): [batch, n_heads, seq_total, head_dim] """ if not (0 <= layer_idx < self.n_layers): raise IndexError(f"layer_idx {layer_idx} fora de range [0, {self.n_layers})") # Validar shapes if new_k.dim() != 4 or new_v.dim() != 4: raise ValueError( f"new_k e new_v devem ser 4D [batch, n_heads, seq, head_dim]; " f"recebido {new_k.shape}" ) # Converter dtype/device new_k = new_k.to(device=self.device, dtype=self.dtype) new_v = new_v.to(device=self.device, dtype=self.dtype) if self.keys[layer_idx] is None: self.keys[layer_idx] = new_k self.values[layer_idx] = new_v else: self.keys[layer_idx] = torch.cat([self.keys[layer_idx], new_k], dim=2) self.values[layer_idx] = torch.cat([self.values[layer_idx], new_v], dim=2) return self.keys[layer_idx], self.values[layer_idx] def evict_sliding(self, max_keep: int) -> None: """ Estratégia sliding window pura: mantém apenas os últimos max_keep tokens. """ for layer_idx in range(self.n_layers): if self.keys[layer_idx] is None: continue seq = self.keys[layer_idx].size(2) if seq > max_keep: start = seq - max_keep self.keys[layer_idx] = self.keys[layer_idx][:, :, start:, :].contiguous() self.values[layer_idx] = self.values[layer_idx][:, :, start:, :].contiguous() def evict_sink_sliding(self, max_keep: int, n_sink: int) -> None: """ Estratégia sink + sliding: mantém os primeiros n_sink + últimos (max_keep - n_sink). """ if n_sink >= max_keep: logger.warning( f"n_sink ({n_sink}) >= max_keep ({max_keep}); usando sliding puro" ) self.evict_sliding(max_keep) return n_sliding = max_keep - n_sink for layer_idx in range(self.n_layers): if self.keys[layer_idx] is None: continue seq = self.keys[layer_idx].size(2) if seq > max_keep: sink_k = self.keys[layer_idx][:, :, :n_sink, :] sink_v = self.values[layer_idx][:, :, :n_sink, :] sliding_k = self.keys[layer_idx][:, :, -n_sliding:, :] sliding_v = self.values[layer_idx][:, :, -n_sliding:, :] self.keys[layer_idx] = torch.cat([sink_k, sliding_k], dim=2).contiguous() self.values[layer_idx] = torch.cat([sink_v, sliding_v], dim=2).contiguous() def reset(self) -> None: """Limpa o cache completamente.""" self.keys = [None] * self.n_layers self.values = [None] * self.n_layers def get_seq_len(self, layer_idx: int = 0) -> int: if self.keys[layer_idx] is None: return 0 return self.keys[layer_idx].size(2) def total_tokens(self) -> int: return max(self.get_seq_len(i) for i in range(self.n_layers)) # ============================================================================ # Context Window Manager # ============================================================================ class ContextWindowManager: """ Gerencia a janela de contexto aplicando a política de evicção configurada. """ def __init__(self, config: ContextWindowConfig): self.config = config self.head_dim = config.head_dim or (config.embed_dim // config.n_heads) if config.embed_dim % config.n_heads != 0: raise ValueError( f"embed_dim ({config.embed_dim}) deve ser divisível por " f"n_heads ({config.n_heads})" ) self.kv_cache: Optional[KVCache] = None # Sequência de tokens brutos (para recomputação se necessário) self.token_history: List[torch.Tensor] = [] # ---------------------------------------------------------------------- # Inicialização # ---------------------------------------------------------------------- def init_cache(self, batch_size: int, device: torch.device, dtype: Optional[torch.dtype] = None) -> KVCache: """Cria novo cache KV para uma sessão de geração.""" dtype = dtype or self.config.dtype or torch.float32 self.kv_cache = KVCache( n_layers=self.config.n_layers, batch_size=batch_size, n_heads=self.config.n_heads, head_dim=self.head_dim, device=device, dtype=dtype, ) self.token_history = [] return self.kv_cache # ---------------------------------------------------------------------- # Adicionar tokens # ---------------------------------------------------------------------- def append_tokens(self, token_ids: torch.Tensor) -> torch.Tensor: """ Adiciona tokens ao histórico e retorna a janela ativa. Args: token_ids: [batch, seq_new] Returns: window: [batch, seq_window] — tokens na janela ativa após evicção """ if token_ids.dim() != 2: raise ValueError(f"token_ids deve ser [batch, seq]; recebido {token_ids.shape}") # Para simplificação, mantemos histórico apenas do batch[0] # (geração autoregressiva tipicamente batch=1) if token_ids.size(0) > 1: logger.warning("ContextWindowManager.append_tokens: batch>1, histórico rastreia apenas batch[0]") self.token_history.append(token_ids) # Concatenar todo o histórico full = torch.cat(self.token_history, dim=1) # Aplicar evicção se necessário return self._evict_tokens(full) def _evict_tokens(self, full_seq: torch.Tensor) -> torch.Tensor: """Aplica política de evicção à sequência completa.""" seq_len = full_seq.size(1) max_keep = self.config.max_window if seq_len <= max_keep: return full_seq strategy = self.config.eviction_strategy if strategy == "sliding": return full_seq[:, -max_keep:] elif strategy == "sink_sliding": n_sink = self.config.n_sink_tokens n_sliding = max_keep - n_sink sink = full_seq[:, :n_sink] sliding = full_seq[:, -n_sliding:] return torch.cat([sink, sliding], dim=1) elif strategy == "recompute": # Manter tudo (atenção recomputada) return full_seq else: logger.warning(f"Estratégia de evicção desconhecida: {strategy}; usando sliding") return full_seq[:, -max_keep:] # ---------------------------------------------------------------------- # Evict cache KV # ---------------------------------------------------------------------- def evict_cache(self) -> None: """Aplica a política de evicção ao cache KV.""" if self.kv_cache is None: return strategy = self.config.eviction_strategy max_keep = self.config.max_window if strategy == "sliding": self.kv_cache.evict_sliding(max_keep) elif strategy == "sink_sliding": self.kv_cache.evict_sink_sliding(max_keep, self.config.n_sink_tokens) elif strategy == "recompute": self.kv_cache.reset() else: self.kv_cache.evict_sliding(max_keep) # ---------------------------------------------------------------------- # Reset # ---------------------------------------------------------------------- def reset(self) -> None: """Reseta todo o estado (cache + histórico).""" if self.kv_cache is not None: self.kv_cache.reset() self.kv_cache = None self.token_history = [] # ---------------------------------------------------------------------- # Info # ---------------------------------------------------------------------- def get_info(self) -> Dict: """Retorna informações sobre o estado atual.""" cache_len = self.kv_cache.total_tokens() if self.kv_cache else 0 history_len = sum(t.size(1) for t in self.token_history) return { "max_window": self.config.max_window, "strategy": self.config.eviction_strategy, "n_sink": self.config.n_sink_tokens, "cache_tokens": cache_len, "history_tokens": history_len, "active": self.kv_cache is not None, } # ============================================================================ # Helpers # ============================================================================ def make_context_window( max_window: int = 512, strategy: str = "sink_sliding", n_sink: int = 4, embed_dim: int = 256, n_heads: int = 4, n_layers: int = 2, device: str = "cpu", ) -> ContextWindowManager: """Factory para criar ContextWindowManager com defaults sensatos.""" config = ContextWindowConfig( max_window=max_window, eviction_strategy=strategy, n_sink_tokens=n_sink, embed_dim=embed_dim, n_heads=n_heads, head_dim=embed_dim // n_heads, n_layers=n_layers, device=device, ) return ContextWindowManager(config) __all__ = [ "ContextWindowConfig", "KVCache", "ContextWindowManager", "make_context_window", ]