Project-Chronos 0.1.1__py3-none-any.whl

chronos/model/hybrid_attention.py

@@ -0,0 +1,237 @@
+"""
+chronos/model/hybrid_attention.py
+
+Hybrid attention for Project Chronos:
+- MLAAttention: Multi-head Latent Attention (DeepSeek-style)
+  Compresses KV cache via low-rank projection: hidden_dim → kv_latent_dim
+  Cache stores latent vectors instead of full KV → ~8-16x smaller KV cache
+
+- SlidingWindowAttention: local attention with fixed window
+  KV cache capped at window_size tokens → O(1) memory per step
+
+- HybridAttentionBlock: alternates MLA (even layers) / SlidingWindow (odd layers)
+  Even layers: MLA for global context with compressed cache
+  Odd layers: SlidingWindow for local patterns with bounded cache
+"""
+import sys
+import chronos.deps  # ensure minimind on sys.path
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from model.model_minimind import RMSNorm, apply_rotary_pos_emb, repeat_kv
+from .config import ChronosConfig
+
+
+# ── MLA: Multi-head Latent Attention ─────────────────────────────
+
+class MLAAttention(nn.Module):
+    """
+    Multi-head Latent Attention (DeepSeek MLA style).
+
+    KV compression:
+        c_kv = W_kv_down(x)          [B, S, kv_latent_dim]   ← cached
+        k, v  = W_kv_up(c_kv)        [B, S, nH, dH]          ← recomputed from cache
+
+    Query decoupled RoPE:
+        q_nope = W_q_nope(x)         [B, S, nH, dH - rope_dim]
+        q_rope = W_q_rope(x)         [B, S, nH, rope_dim]     ← RoPE applied here
+        k_rope = W_k_rope(c_kv)      [B, S, nH, rope_dim]     ← RoPE applied here
+
+    Cache stores only c_kv (latent) + k_rope → much smaller than full KV.
+    """
+
+    def __init__(self, config: ChronosConfig):
+        super().__init__()
+        self.n_heads = config.num_attention_heads
+        self.n_kv_heads = config.num_key_value_heads or config.num_attention_heads
+        self.head_dim = config.head_dim
+        self.n_rep = self.n_heads // self.n_kv_heads
+        self.dropout = config.dropout
+
+        # KV latent dimension: much smaller than n_kv_heads * head_dim
+        self.kv_latent_dim = getattr(config, 'kv_latent_dim',
+                                     max(64, self.n_kv_heads * self.head_dim // 8))
+        # RoPE dimension per head (subset of head_dim); must leave room for nope
+        self.rope_dim = min(getattr(config, 'rope_dim', min(64, self.head_dim)), self.head_dim - 1)
+        self.rope_dim = max(1, self.rope_dim)  # at least 1
+        self.nope_dim = self.head_dim - self.rope_dim
+
+        H = config.hidden_size
+
+        # Query projections
+        self.q_nope_proj = nn.Linear(H, self.n_heads * self.nope_dim, bias=False)
+        self.q_rope_proj = nn.Linear(H, self.n_heads * self.rope_dim, bias=False)
+        self.q_nope_norm = RMSNorm(self.nope_dim, eps=config.rms_norm_eps)
+        self.q_rope_norm = RMSNorm(self.rope_dim, eps=config.rms_norm_eps)
+
+        # KV down-projection (what gets cached)
+        self.kv_down_proj = nn.Linear(H, self.kv_latent_dim, bias=False)
+        self.kv_down_norm = RMSNorm(self.kv_latent_dim, eps=config.rms_norm_eps)
+
+        # KV up-projections (recomputed from cache at inference)
+        self.k_nope_proj = nn.Linear(self.kv_latent_dim, self.n_kv_heads * self.nope_dim, bias=False)
+        self.k_rope_proj = nn.Linear(self.kv_latent_dim, self.n_kv_heads * self.rope_dim, bias=False)
+        self.v_proj = nn.Linear(self.kv_latent_dim, self.n_kv_heads * self.head_dim, bias=False)
+
+        self.o_proj = nn.Linear(self.n_heads * self.head_dim, H, bias=False)
+        self.attn_dropout = nn.Dropout(self.dropout)
+        self.resid_dropout = nn.Dropout(self.dropout)
+        self.flash = hasattr(F, 'scaled_dot_product_attention') and config.flash_attn
+
+    def forward(self, x, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None):
+        B, S, _ = x.shape
+        cos, sin = position_embeddings
+
+        # Query
+        q_nope = self.q_nope_norm(
+            self.q_nope_proj(x).view(B, S, self.n_heads, self.nope_dim)
+        )
+        q_rope = self.q_rope_norm(
+            self.q_rope_proj(x).view(B, S, self.n_heads, self.rope_dim)
+        )
+        # apply_rotary_pos_emb expects [B, S, nH, d] (NOT transposed)
+        # cos/sin are full-length buffers; slice to current step positions
+        past_len = past_key_value[0].shape[1] if past_key_value is not None else 0
+        cos_r = cos[past_len:past_len + S, :self.rope_dim]   # [S, rope_dim]
+        sin_r = sin[past_len:past_len + S, :self.rope_dim]
+        q_rope, _ = apply_rotary_pos_emb(q_rope, q_rope, cos_r, sin_r)  # [B, S, nH, rope_dim]
+        q = torch.cat([q_nope, q_rope], dim=-1)  # [B, S, nH, head_dim]
+
+        # KV latent (this is what we cache)
+        c_kv = self.kv_down_norm(self.kv_down_proj(x))  # [B, S, kv_latent_dim]
+
+        if past_key_value is not None:
+            past_c_kv = past_key_value[0]
+            c_kv_full = torch.cat([past_c_kv, c_kv], dim=1)
+        else:
+            c_kv_full = c_kv
+
+        full_len = c_kv_full.shape[1]
+
+        # Expand KV from latent
+        k_nope = self.k_nope_proj(c_kv_full).view(B, full_len, self.n_kv_heads, self.nope_dim)
+        k_rope_raw = self.k_rope_proj(c_kv_full).view(B, full_len, self.n_kv_heads, self.rope_dim)
+        v = self.v_proj(c_kv_full).view(B, full_len, self.n_kv_heads, self.head_dim)
+
+        # RoPE on keys over full cached sequence — [B, T, nKV, rope_dim]
+        cos_full = cos[:full_len, :self.rope_dim]
+        sin_full = sin[:full_len, :self.rope_dim]
+        k_rope, _ = apply_rotary_pos_emb(k_rope_raw, k_rope_raw, cos_full, sin_full)
+        k = torch.cat([k_nope, k_rope], dim=-1)  # [B, T, nKV, head_dim]
+
+        past_kv = (c_kv_full,) if use_cache else None
+
+        # Attention
+        q = q.transpose(1, 2)                          # [B, nH, S, dH]
+        k = repeat_kv(k, self.n_rep).transpose(1, 2)   # [B, nH, T, dH]
+        v = repeat_kv(v, self.n_rep).transpose(1, 2)   # [B, nH, T, dH]
+
+        T = k.shape[2]
+        if self.flash and S > 1 and attention_mask is None:
+            out = F.scaled_dot_product_attention(
+                q, k, v,
+                dropout_p=self.dropout if self.training else 0.0,
+                is_causal=True,
+            )
+        else:
+            scores = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim)
+            scores[:, :, :, -S:] += torch.full((S, S), float('-inf'), device=scores.device).triu(1)
+            if attention_mask is not None:
+                scores += (1.0 - attention_mask.unsqueeze(1).unsqueeze(2)) * -1e9
+            out = self.attn_dropout(F.softmax(scores.float(), dim=-1).type_as(q)) @ v
+
+        out = out.transpose(1, 2).reshape(B, S, -1)
+        return self.resid_dropout(self.o_proj(out)), past_kv
+
+
+# ── Sliding Window Attention ──────────────────────────────────────
+
+class SlidingWindowAttention(nn.Module):
+    """
+    Local sliding window attention. KV cache is capped at window_size tokens,
+    preventing unbounded growth in long-context generation.
+
+    During prefill: full causal attention within each window.
+    During decode:  KV cache evicts oldest tokens beyond window_size.
+    """
+
+    def __init__(self, config: ChronosConfig):
+        super().__init__()
+        self.n_heads = config.num_attention_heads
+        self.n_kv_heads = config.num_key_value_heads or config.num_attention_heads
+        self.head_dim = config.head_dim
+        self.n_rep = self.n_heads // self.n_kv_heads
+        self.dropout = config.dropout
+        self.window_size = getattr(config, 'sliding_window_size', 2048)
+
+        H = config.hidden_size
+        self.q_proj = nn.Linear(H, self.n_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(H, self.n_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(H, self.n_kv_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.n_heads * self.head_dim, H, bias=False)
+        self.q_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.k_norm = RMSNorm(self.head_dim, eps=config.rms_norm_eps)
+        self.attn_dropout = nn.Dropout(self.dropout)
+        self.resid_dropout = nn.Dropout(self.dropout)
+        self.flash = hasattr(F, 'scaled_dot_product_attention') and config.flash_attn
+
+    def forward(self, x, position_embeddings, past_key_value=None, use_cache=False, attention_mask=None):
+        B, S, _ = x.shape
+        cos, sin = position_embeddings
+        # cos/sin are full-length buffers; slice to current positions
+        past_len = past_key_value[0].shape[1] if past_key_value is not None else 0
+        cos_s = cos[past_len:past_len + S]
+        sin_s = sin[past_len:past_len + S]
+
+        xq = self.q_norm(self.q_proj(x).view(B, S, self.n_heads, self.head_dim))
+        xk = self.k_norm(self.k_proj(x).view(B, S, self.n_kv_heads, self.head_dim))
+        xv = self.v_proj(x).view(B, S, self.n_kv_heads, self.head_dim)
+
+        xq, xk = apply_rotary_pos_emb(xq, xk, cos_s, sin_s)  # [B, S, nH/nKV, d]
+
+        if past_key_value is not None:
+            xk = torch.cat([past_key_value[0], xk], dim=1)
+            xv = torch.cat([past_key_value[1], xv], dim=1)
+
+        # Evict tokens beyond window_size (sliding window)
+        if xk.shape[1] > self.window_size:
+            xk = xk[:, -self.window_size:, :, :]
+            xv = xv[:, -self.window_size:, :, :]
+
+        past_kv = (xk, xv) if use_cache else None
+
+        xq = xq.transpose(1, 2)
+        xk = repeat_kv(xk, self.n_rep).transpose(1, 2)
+        xv = repeat_kv(xv, self.n_rep).transpose(1, 2)
+
+        T = xk.shape[2]
+        if self.flash and S > 1 and attention_mask is None:
+            out = F.scaled_dot_product_attention(
+                xq, xk, xv,
+                dropout_p=self.dropout if self.training else 0.0,
+                is_causal=(T == S),
+            )
+        else:
+            scores = (xq @ xk.transpose(-2, -1)) / math.sqrt(self.head_dim)
+            if T == S:
+                scores[:, :, :, -S:] += torch.full((S, S), float('-inf'), device=scores.device).triu(1)
+            if attention_mask is not None:
+                scores += (1.0 - attention_mask[:, :, :, -T:].unsqueeze(1)) * -1e9
+            out = self.attn_dropout(F.softmax(scores.float(), dim=-1).type_as(xq)) @ xv
+
+        out = out.transpose(1, 2).reshape(B, S, -1)
+        return self.resid_dropout(self.o_proj(out)), past_kv
+
+
+def make_attention(layer_id: int, config: ChronosConfig) -> nn.Module:
+    """
+    Layer assignment:
+      even layers → MLAAttention  (global context, compressed KV cache)
+      odd  layers → SlidingWindowAttention (local patterns, bounded KV cache)
+    """
+    if layer_id % 2 == 0:
+        return MLAAttention(config)
+    return SlidingWindowAttention(config)

chronos/model/lookahead_router.py

@@ -0,0 +1,55 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .config import ChronosConfig
+
+
+class LookaheadRouter(nn.Module):
+    """
+    Global lookahead routing predictor inserted after the first transformer block.
+
+    Input:  hidden_states [B, S, H]  (output of block 0)
+    Output: expert_probs  [B, S, lookahead_steps+1, num_experts]
+              dim-2 index 0 = current step t
+              dim-2 index 1 = predicted t+1
+              dim-2 index 2 = predicted t+2  (if lookahead_steps >= 2)
+
+    During training: output drives temporal_locality_loss and lookahead accuracy.
+    During inference: output drives async prefetch for future expert weights.
+    """
+
+    def __init__(self, config: ChronosConfig):
+        super().__init__()
+        self.lookahead_steps = config.lookahead_steps
+        self.num_experts = config.num_experts
+        H = config.hidden_size
+        K = config.lookahead_steps + 1  # current + future steps
+
+        # Lightweight 2-layer MLP: H -> H//4 -> E*K
+        self.proj = nn.Sequential(
+            nn.Linear(H, H // 4, bias=False),
+            nn.SiLU(),
+            nn.Linear(H // 4, self.num_experts * K, bias=False),
+        )
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        # hidden_states: [B, S, H]
+        B, S, H = hidden_states.shape
+        K = self.lookahead_steps + 1
+        logits = self.proj(hidden_states)                    # [B, S, E*K]
+        logits = logits.view(B, S, K, self.num_experts)      # [B, S, K, E]
+        return F.softmax(logits, dim=-1)                     # [B, S, K, E]
+
+    def predict_next_experts(
+        self,
+        hidden_states: torch.Tensor,
+        top_k: int = 1,
+    ) -> torch.Tensor:
+        """
+        Inference helper: returns top-k expert indices for each future step.
+        Returns: [B, S, lookahead_steps, top_k]
+        """
+        probs = self.forward(hidden_states)          # [B, S, K, E]
+        future = probs[:, :, 1:, :]                  # [B, S, lookahead_steps, E]
+        return torch.topk(future, k=top_k, dim=-1).indices

chronos/model/model_chronos.py

@@ -0,0 +1,271 @@
+import sys
+import chronos.deps  # ensure minimind on sys.path
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel, GenerationMixin
+from transformers.modeling_outputs import MoeCausalLMOutputWithPast
+
+from model.model_minimind import (
+    RMSNorm, precompute_freqs_cis, Attention, FeedForward,
+)
+from .config import ChronosConfig
+from .lookahead_router import LookaheadRouter
+from .moe_chronos import ChronosMOEFeedForward
+from .hybrid_attention import make_attention
+
+
+class ChronosBlock(nn.Module):
+    """
+    Transformer block with hybrid attention:
+      even layers → MLAAttention  (compressed KV cache, global context)
+      odd  layers → SlidingWindowAttention (bounded KV cache, local context)
+    """
+
+    def __init__(self, layer_id: int, config: ChronosConfig):
+        super().__init__()
+        self.layer_id = layer_id
+        if getattr(config, 'use_hybrid_attention', True):
+            self.self_attn = make_attention(layer_id, config)
+        else:
+            self.self_attn = Attention(config)
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.mlp = ChronosMOEFeedForward(config) if config.use_moe else FeedForward(config)
+
+    def forward(
+        self,
+        hidden_states,
+        position_embeddings,
+        past_key_value=None,
+        use_cache=False,
+        attention_mask=None,
+        available_expert_mask=None,
+    ):
+        residual = hidden_states
+        hidden_states, present_key_value = self.self_attn(
+            self.input_layernorm(hidden_states),
+            position_embeddings,
+            past_key_value,
+            use_cache,
+            attention_mask,
+        )
+        hidden_states = hidden_states + residual
+        mlp_in = self.post_attention_layernorm(hidden_states)
+        if isinstance(self.mlp, ChronosMOEFeedForward):
+            hidden_states = hidden_states + self.mlp(mlp_in, available_expert_mask)
+        else:
+            hidden_states = hidden_states + self.mlp(mlp_in)
+        return hidden_states, present_key_value
+
+
+class ChronosModel(nn.Module):
+    """
+    Backbone: identical to MiniMindModel but:
+    - Uses ChronosBlock (with ChronosMOEFeedForward)
+    - Inserts LookaheadRouter after block 0
+    - Exposes router_probs_seq for temporal loss computation
+    """
+
+    def __init__(self, config: ChronosConfig):
+        super().__init__()
+        self.config = config
+        self.vocab_size = config.vocab_size
+        self.num_hidden_layers = config.num_hidden_layers
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.dropout = nn.Dropout(config.dropout)
+        self.layers = nn.ModuleList([
+            ChronosBlock(l, config) for l in range(self.num_hidden_layers)
+        ])
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.lookahead_router = LookaheadRouter(config)
+        freqs_cos, freqs_sin = precompute_freqs_cis(
+            dim=config.head_dim,
+            end=config.max_position_embeddings,
+            rope_base=config.rope_theta,
+            rope_scaling=config.rope_scaling,
+        )
+        self.register_buffer("freqs_cos", freqs_cos, persistent=False)
+        self.register_buffer("freqs_sin", freqs_sin, persistent=False)
+
+    def forward(
+        self,
+        input_ids,
+        attention_mask=None,
+        past_key_values=None,
+        use_cache=False,
+        available_expert_masks=None,  # list[Tensor|None] per layer
+        **kwargs,
+    ):
+        B, S = input_ids.shape
+        if hasattr(past_key_values, 'layers'):
+            past_key_values = None
+        past_key_values = past_key_values or [None] * len(self.layers)
+        # MLA caches only c_kv tuple; SlidingWindow caches (k, v) tuple
+        # Determine start_pos from first non-None past_kv
+        start_pos = 0
+        for pkv in past_key_values:
+            if pkv is not None:
+                start_pos = pkv[0].shape[1]
+                break
+
+        hidden_states = self.dropout(self.embed_tokens(input_ids))
+        # Pass full freqs so hybrid attention layers can slice as needed
+        position_embeddings = (
+            self.freqs_cos,
+            self.freqs_sin,
+        )
+
+        presents = []
+        lookahead_probs = None  # [B, S, K, E] from block 0
+
+        for i, (layer, past_kv) in enumerate(zip(self.layers, past_key_values)):
+            mask = available_expert_masks[i] if available_expert_masks else None
+            hidden_states, present = layer(
+                hidden_states, position_embeddings,
+                past_key_value=past_kv,
+                use_cache=use_cache,
+                attention_mask=attention_mask,
+                available_expert_mask=mask,
+            )
+            presents.append(present)
+            # Insert LookaheadRouter after block 0
+            if i == 0:
+                lookahead_probs = self.lookahead_router(hidden_states)
+
+        hidden_states = self.norm(hidden_states)
+        aux_loss = sum(
+            [l.mlp.aux_loss for l in self.layers if isinstance(l.mlp, ChronosMOEFeedForward)],
+            hidden_states.new_zeros(1).squeeze(),
+        )
+        return hidden_states, presents, aux_loss, lookahead_probs
+
+
+class ChronosForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = ChronosConfig
+    # Map of tied-weight keys (lm_head.weight ↔ model.embed_tokens.weight).
+    # Newer Transformers expects a dict; older versions accept a list. We
+    # use a dict to satisfy the newest API.
+    _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
+
+    @property
+    def all_tied_weights_keys(self) -> dict:
+        """Compat shim: newer Transformers expects this property to exist
+        on every PreTrainedModel; older base classes don't define it."""
+        v = self._tied_weights_keys
+        if isinstance(v, dict):
+            return v
+        if isinstance(v, (list, tuple, set)):
+            return {k: None for k in v}
+        return {}
+
+    def __init__(self, config: ChronosConfig = None):
+        self.config = config or ChronosConfig()
+        super().__init__(self.config)
+        self.model = ChronosModel(self.config)
+        self.lm_head = nn.Linear(self.config.hidden_size, self.config.vocab_size, bias=False)
+        if self.config.tie_word_embeddings:
+            self.model.embed_tokens.weight = self.lm_head.weight
+
+    def forward(
+        self,
+        input_ids,
+        attention_mask=None,
+        past_key_values=None,
+        use_cache=False,
+        logits_to_keep=0,
+        labels=None,
+        available_expert_masks=None,
+        **kwargs,
+    ):
+        hidden_states, past_key_values, aux_loss, lookahead_probs = self.model(
+            input_ids, attention_mask, past_key_values, use_cache,
+            available_expert_masks=available_expert_masks,
+        )
+        slice_indices = slice(-logits_to_keep, None) if logits_to_keep > 0 else slice(None)
+        logits = self.lm_head(hidden_states[:, slice_indices, :])
+
+        loss = None
+        if labels is not None:
+            x = logits[..., :-1, :].contiguous()
+            y = labels[..., 1:].contiguous()
+            loss = F.cross_entropy(x.view(-1, x.size(-1)), y.view(-1), ignore_index=-100)
+
+        return MoeCausalLMOutputWithPast(
+            loss=loss,
+            aux_loss=aux_loss,
+            logits=logits,
+            past_key_values=past_key_values,
+            hidden_states=hidden_states,
+        ), lookahead_probs
+
+    @torch.inference_mode()
+    def generate(
+        self,
+        inputs=None,
+        attention_mask=None,
+        max_new_tokens=512,
+        temperature=0.85,
+        top_p=0.85,
+        top_k=50,
+        eos_token_id=2,
+        streamer=None,
+        use_cache=True,
+        do_sample=True,
+        repetition_penalty=1.0,
+        available_expert_masks=None,
+        **kwargs,
+    ):
+        input_ids = kwargs.pop("input_ids", inputs)
+        past_key_values = kwargs.pop("past_key_values", None)
+        finished = torch.zeros(input_ids.shape[0], dtype=torch.bool, device=input_ids.device)
+        if streamer:
+            streamer.put(input_ids.cpu())
+
+        for _ in range(max_new_tokens):
+            past_len = past_key_values[0][0].shape[1] if past_key_values else 0
+            outputs, _ = self.forward(
+                input_ids[:, past_len:], attention_mask, past_key_values,
+                use_cache=use_cache,
+                available_expert_masks=available_expert_masks,
+            )
+            if attention_mask is not None:
+                attention_mask = torch.cat(
+                    [attention_mask, attention_mask.new_ones(attention_mask.shape[0], 1)], -1
+                )
+            logits = outputs.logits[:, -1, :] / temperature
+            if repetition_penalty != 1.0:
+                for i in range(input_ids.shape[0]):
+                    logits[i, torch.unique(input_ids[i])] /= repetition_penalty
+            if top_k > 0:
+                logits[logits < torch.topk(logits, top_k)[0][..., -1, None]] = -float('inf')
+            if top_p < 1.0:
+                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+                mask = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1) > top_p
+                mask[..., 1:] = mask[..., :-1].clone()
+                mask[..., 0] = 0
+                logits[mask.scatter(1, sorted_indices, mask)] = -float('inf')
+            next_token = (
+                torch.multinomial(torch.softmax(logits, dim=-1), num_samples=1)
+                if do_sample else torch.argmax(logits, dim=-1, keepdim=True)
+            )
+            if eos_token_id is not None:
+                next_token = torch.where(
+                    finished.unsqueeze(-1),
+                    next_token.new_full((next_token.shape[0], 1), eos_token_id),
+                    next_token,
+                )
+            input_ids = torch.cat([input_ids, next_token], dim=-1)
+            past_key_values = outputs.past_key_values if use_cache else None
+            if streamer:
+                streamer.put(next_token.cpu())
+            if eos_token_id is not None:
+                finished |= next_token.squeeze(-1).eq(eos_token_id)
+                if finished.all():
+                    break
+
+        if streamer:
+            streamer.end()
+        return input_ids