oicio/memory/reattention.py

"""
OICIO Memory Fabric: ReAttention - Training-Free Infinite Context with Finite Attention Scope
Credits: deepRcurs Labs, @deeprcurs / Mzed Imamkh @mzedimamkh

Berdasarkan:
- ReAttention paper 2407.15176v3
- 3 syarat infinite context: pos emb not OOD, stable entropy, effective awareness

Core:
1. Position-agnostic top-k attention BEFORE position-aware attention
2. q_t * K_middle^T tanpa RoPE untuk cari critical info
3. Concat [K_global 32 + K_select 127*32 + K_local 4096] = 8192 max, baru kasih RoPE
"""

import numpy as np
import torch
import torch.nn.functional as F
from typing import Tuple

class ReAttention:
    """
    ReAttention: finite attention scope, infinite context
    """
    def __init__(self,
                 global_tokens: int = 32,
                 local_tokens: int = 4096,
                 select_span: int = 32,
                 top_k: int = 4,
                 top_k_prime: int = 127):
        self.global_tokens = global_tokens
        self.local_tokens = local_tokens
        self.select_span = select_span
        self.top_k = top_k
        self.top_k_prime = top_k_prime

        # Max attention scope = global + local + k' * span
        self.max_scope = global_tokens + local_tokens + top_k_prime * select_span
        print(f"[ReAttention] Max attention scope: {self.max_scope} (global={global_tokens} + local={local_tokens} + {top_k_prime}*{select_span})")

    def split_cache(self, kv_cache: np.ndarray):
        """
        Split KV cache into global, middle, local
        kv_cache: [seq_len, dim]
        """
        seq_len = kv_cache.shape[0]

        if seq_len <= self.global_tokens + self.local_tokens:
            # Not enough to split, return all as local
            return kv_cache[:0], kv_cache, kv_cache[:0]

        global_part = kv_cache[:self.global_tokens]
        local_part = kv_cache[-self.local_tokens:]
        middle_part = kv_cache[self.global_tokens:-self.local_tokens]

        return global_part, middle_part, local_part

    def position_agnostic_selection(self, 
                                   query: np.ndarray, 
                                   middle_k: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        Position-agnostic top-k selection
        query: [dim] or [1, dim]
        middle_k: [middle_len, dim]
        Returns: indices of top-k' spans and selected K/V

        Real paper uses Triton kernel fused top-k
        POC uses numpy
        """
        if query.ndim == 1:
            query = query[None, :]

        # Dot product without position embedding
        # q_t * K_middle^T
        scores = query @ middle_k.T  # [1, middle_len]
        scores = scores[0]  # [middle_len]

        # Top-k' selection (k' = 127)
        # But we also want to consider multiple heads/queries voting
        # For POC, simple top-k

        # To ensure semantic coherence, not only top-k' elements but also m neighboring entries
        # Overlapping parts deduplicated
        top_indices = np.argsort(scores)[-self.top_k_prime*2:][::-1]  # get more, then dedup spans

        # Expand each index to span
        selected_indices = set()
        for idx in top_indices:
            start = max(0, idx - self.select_span//2)
            end = min(len(middle_k), start + self.select_span)
            for j in range(start, end):
                selected_indices.add(j)
            if len(selected_indices) >= self.top_k_prime * self.select_span:
                break

        selected_indices = sorted(list(selected_indices))[:self.top_k_prime * self.select_span]

        return np.array(selected_indices), scores[selected_indices] if len(selected_indices) > 0 else np.array([])

    def reconstruct_cache(self,
                         global_k: np.ndarray,
                         select_k: np.ndarray,
                         local_k: np.ndarray,
                         global_v: np.ndarray = None,
                         select_v: np.ndarray = None,
                         local_v: np.ndarray = None):
        """
        Concatenate selected segments between global and local
        Preserves relative order while ignoring absolute distance
        Then apply RoPE sequentially (simulated)
        """
        # Concatenate K
        if len(select_k) > 0:
            k_concat = np.concatenate([global_k, select_k, local_k], axis=0)
        else:
            k_concat = np.concatenate([global_k, local_k], axis=0)

        # V similarly
        if global_v is not None:
            if select_v is not None and len(select_v) > 0:
                v_concat = np.concatenate([global_v, select_v, local_v], axis=0)
            else:
                v_concat = np.concatenate([global_v, local_v], axis=0)
        else:
            v_concat = None

        # Apply position embedding sequentially (RoPE simulation)
        # In real ReAttention, PE is separated from KV cache and performed AFTER selection
        # This ensures PE never OOD because concat length <= pretrain window
        # For POC, we just return concat, PE would be applied in attention

        return k_concat, v_concat

    def forward(self, query: np.ndarray, kv_cache: np.ndarray, v_cache: np.ndarray = None):
        """
        Full ReAttention forward
        query: [dim] current query
        kv_cache: [seq_len, dim] full cache
        v_cache: [seq_len, dim] optional V cache

        Returns: selected K,V for attention
        """
        global_k, middle_k, local_k = self.split_cache(kv_cache)

        if v_cache is not None:
            global_v, middle_v, local_v = self.split_cache(v_cache)
        else:
            global_v, middle_v, local_v = None, None, None
            middle_v = middle_k  # for simplicity

        if len(middle_k) == 0:
            # No middle, just global + local
            k_concat = np.concatenate([global_k, local_k], axis=0) if len(global_k) > 0 else local_k
            v_concat = np.concatenate([global_v, local_v], axis=0) if global_v is not None and len(global_v) > 0 else local_v
            return k_concat, v_concat, np.array([])

        # Position-agnostic selection
        select_indices, select_scores = self.position_agnostic_selection(query, middle_k)

        if len(select_indices) > 0:
            select_k = middle_k[select_indices]
            select_v = middle_v[select_indices] if middle_v is not None else select_k
        else:
            select_k = np.zeros((0, kv_cache.shape[1]))
            select_v = np.zeros((0, kv_cache.shape[1])) if v_cache is not None else None

        # Reconstruct
        k_final, v_final = self.reconstruct_cache(global_k, select_k, local_k, global_v, select_v, local_v)

        return k_final, v_final, select_indices

    def attention(self, query: np.ndarray, k_cache: np.ndarray, v_cache: np.ndarray):
        """
        Self-attention with selected cache
        query: [dim]
        k_cache, v_cache: selected caches [selected_len, dim]
        """
        # Apply RoPE sequentially (simulate as no-op for POC, but ensure length within window)
        assert len(k_cache) <= self.max_scope, f"Cache {len(k_cache)} exceeds max scope {self.max_scope}, would be OOD!"

        # Standard attention
        scores = query @ k_cache.T  # [selected_len]
        scores = scores / np.sqrt(query.shape[0])
        attn_weights = np.exp(scores - np.max(scores))
        attn_weights = attn_weights / np.sum(attn_weights)

        # Output
        out = attn_weights @ v_cache  # [dim]

        return out, attn_weights

# Demo
if __name__ == "__main__":
    print("=== ReAttention POC ===")
    dim = 64
    seq_len = 100000  # 100K context

    # Simulate KV cache
    kv_cache = np.random.randn(seq_len, dim).astype(np.float32)
    v_cache = np.random.randn(seq_len, dim).astype(np.float32)

    # Current query
    query = np.random.randn(dim).astype(np.float32)

    reatt = ReAttention(global_tokens=32, local_tokens=128, select_span=32, top_k_prime=10)  # small for POC

    k_final, v_final, indices = reatt.forward(query, kv_cache, v_cache)

    print(f"Original cache: {seq_len}")
    print(f"Selected cache: {len(k_final)} (global 32 + select {len(indices)} + local 128)")
    print(f"Compression: {seq_len} -> {len(k_final)} = {seq_len/len(k_final):.1f}x")
    print(f"Within max scope {reatt.max_scope}? {len(k_final) <= reatt.max_scope}")

    out, weights = reatt.attention(query, k_final, v_final)
    print(f"Attention output shape: {out.shape}")
    print(f"Attention entropy: {-np.sum(weights * np.log(weights + 1e-8)):.3f} (should be stable, not grow with seq_len)")