BeaconKV 2026: Arsitektur Kompresi KV Cache Berpemandu Beacon Queries, Thought Revisiting Tokens (TRT), dan Akselerasi Inferensi Large Reasoning Models Tanpa Retraining

"""
BeaconKV Production Core Engine: Continual FPS & Dual-Scoring KV Cache Compression
Berdasarkan Arsitektur Resmi arXiv:2609.04971 (Kim et al., AIHA Lab)
Mendukung FlashAttention-2, Group Max Aggregation, dan Pre-RoPE Continual Sampling.
"""

import torch
import torch.nn.functional as F
from typing import Optional, Tuple

class BeaconKVCacheCompressor:
    """
    Manajer kompresi KV Cache berpemandu Beacon Queries.
    Mempertahankan query perwakilan klaster embedding menggunakan Continual Furthest Point Sampling (FPS)
    dan mengeksekusi eviksi terarah pada KV cache tanpa membutuhkan training ulang.
    """
    def __init__(
        self,
        max_batch_size: int,
        max_input_len: int,
        decode_max_budget: int,
        decode_min_budget: int,
        num_recent_queries: int,
        num_beacon_queries: int,
        num_attention_heads: int,
        head_dim: int,
        device: torch.device,
        dtype: torch.dtype = torch.bfloat16,
    ):
        self.max_batch_size = max_batch_size
        self.max_input_len = max_input_len
        self.decode_max_budget = decode_max_budget
        self.decode_min_budget = decode_min_budget
        self.num_recent_queries = num_recent_queries
        self.num_beacon_queries = num_beacon_queries
        self.num_max_queries = num_recent_queries + num_beacon_queries
        self.num_attention_heads = num_attention_heads
        self.head_dim = head_dim
        self.device = device
        self.dtype = dtype

        # Pre-RoPE query buffer shape: [B, H, Q_max, D]
        query_cache_shape = (
            self.max_batch_size,
            self.num_attention_heads,
            self.num_max_queries,
            self.head_dim,
        )
        self.query_cache = torch.zeros(query_cache_shape, dtype=self.dtype, device=self.device)
        self.query_counter = 0

    def update_pre_rope_query(
        self,
        query_states: torch.Tensor,
        current_kv_len: int,
        is_prefill: bool,
        attention_mask: Optional[torch.Tensor] = None,
    ) -> None:
        """
        Memperbarui query cache pre-RoPE secara continual.
        Pada tahap prefill: inisialisasi beacon queries dari prompt.
        Pada tahap decode: akumulasi query recent dan continual resampling saat buffer penuh.
        """
        if is_prefill:
            # Inisialisasi beacon queries awal via FPS dari prompt prefill
            self.query_cache[:, :, :self.num_beacon_queries, :] = self._farthest_point_sampling(
                query_states, self.num_beacon_queries, attention_mask
            )
            self.query_counter = self.num_beacon_queries
            return

        # Fase decoding: periksa apakah buffer melebihi kapasitas
        if (current_kv_len - self.max_input_len) <= self.decode_max_budget:
            if self.query_counter == self.num_max_queries:
                # Refresh klaster beacon queries secara on-the-fly (Continual FPS)
                self.query_cache[:, :, :self.num_beacon_queries, :] = self._farthest_point_sampling(
                    self.query_cache[:, :, :self.query_counter, :],
                    self.num_beacon_queries,
                    attention_mask=None,
                )
                self.query_counter = self.num_beacon_queries

            # Sisipkan decode query terbaru ke slot recent
            self.query_cache[:, :, self.query_counter, :] = query_states.squeeze(2)
            self.query_counter += 1

    def _farthest_point_sampling(
        self,
        query_states: torch.Tensor,
        num_samples: int,
        attention_mask: Optional[torch.Tensor] = None,
    ) -> torch.Tensor:
        """
        Continual Furthest Point Sampling (FPS) berbasis Cosine Similarity.
        Memilih titik representatif yang memaksimalkan jarak sudut antar vektor kueri.
        Kompleksitas: O(Q^2) dengan precomputed Gram Matrix.
        """
        batch_size, num_heads, num_queries, head_dim = query_states.shape
        if num_queries <= num_samples:
            return query_states

        num_groups = batch_size * num_heads
        q_flat = query_states.reshape(num_groups, num_queries, head_dim)
        q_norm = F.normalize(q_flat, p=2, dim=-1)

        # Hitung pairwise similarity matrix: [num_groups, Q, Q]
        sim_matrix = torch.bmm(q_norm, q_norm.transpose(-2, -1))
        
        # Cari titik awal: kueri dengan rata-rata kemiripan terendah terhadap kueri lainnya
        mean_sim = sim_matrix.mean(dim=2)
        start_idx = mean_sim.argmin(dim=-1)

        selected_indices = torch.empty(num_groups, num_samples, dtype=torch.long, device=self.device)
        selected_indices[:, 0] = start_idx

        # Inisialisasi pelacak kemiripan maksimum (similairest distance)
        g_idx = torch.arange(num_groups, device=self.device)
        max_sim = sim_matrix[g_idx, start_idx] # [num_groups, Q]

        for i in range(1, num_samples):
            # Titik berikutnya adalah titik yang memiliki kemiripan maksimum paling rendah (furthest point)
            next_idx = max_sim.argmin(dim=-1)
            selected_indices[:, i] = next_idx
            
            # Perbarui jarak batas
            curr_sim = sim_matrix[g_idx, next_idx]
            max_sim = torch.maximum(max_sim, curr_sim)

        # Kumpulkan tensor kueri terpilih
        gather_idx = selected_indices.unsqueeze(-1).expand(-1, -1, head_dim)
        sampled_q = torch.gather(q_flat, dim=1, index=gather_idx)
        return sampled_q.reshape(batch_size, num_heads, num_samples, head_dim)

    def compress_kv_cache(
        self,
        key_states: torch.Tensor,
        value_states: torch.Tensor,
        min_token_budget: int,
        kv_prefix_len: int,
        kv_recent_len: int,
        post_rope_beacon_and_recent_queries: torch.Tensor,
    ) -> Tuple[torch.Tensor, torch.Tensor]:
        """
        Menjalankan seleksi retensi KV cache:
        1. Menghitung matriks atensi antara queries (beacons + recent) dan key states.
        2. Menerapkan group_max head aggregation untuk mengakomodasi Grouped Query Attention (GQA).
        3. Menjaga prefix prompt asli dan local recent sliding window utuh.
        4. Mengeliminasi token intermediate yang tidak pernah dikunjungi kembali oleh beacon queries.
        """
        batch_size, num_kv_heads, seq_len, head_dim = key_states.shape
        remaining_budget = min_token_budget - (kv_prefix_len + kv_recent_len)
        
        if remaining_budget <= 0:
            return key_states, value_states

        # Hitung skor atensi teragregasi
        # [B, H_q, Q_active, S] @ [B, H_kv, S, D] -> skor perhatian
        scores = torch.einsum('bhqd,bksd->bhqs', post_rope_beacon_and_recent_queries, key_states)
        scores = scores / (head_dim ** 0.5)

        # Aggregasi Group Max untuk query groups pada GQA
        num_query_heads = post_rope_beacon_and_recent_queries.shape[1]
        group_size = num_query_heads // num_kv_heads
        scores = scores.view(batch_size, num_kv_heads, group_size, -1, seq_len)
        scores = scores.max(dim=2).values # Reduksi group_size
        scores = scores.max(dim=2).values # Reduksi kueri (beacon + recent max)

        # Ambil hanya area kompresi (antara prefix dan recent window)
        eviction_candidate_scores = scores[:, :, kv_prefix_len:seq_len - kv_recent_len]
        
        # Top-K pemilihan token esensial berdasarkan skor tertinggi
        topk_indices = eviction_candidate_scores.topk(remaining_budget, dim=-1, largest=True).indices
        sorted_topk_indices = topk_indices.sort(dim=-1).values + kv_prefix_len

        # Susun indeks gabungan: [Prefix, Top-K Preserved, Recent Window]
        prefix_idx = torch.arange(kv_prefix_len, device=self.device).expand(batch_size, num_kv_heads, -1)
        recent_idx = torch.arange(seq_len - kv_recent_len, seq_len, device=self.device).expand(batch_size, num_kv_heads, -1)
        final_indices = torch.cat([prefix_idx, sorted_topk_indices, recent_idx], dim=-1)

        # Ekstrak KV cache terkompresi
        gather_indices = final_indices.unsqueeze(-1).expand(-1, -1, -1, head_dim)
        compressed_k = torch.gather(key_states, dim=2, index=gather_indices)
        compressed_v = torch.gather(value_states, dim=2, index=gather_indices)

        return compressed_k, compressed_v