Upload python/utils/chamber_index.py with huggingface_hub

python/utils/chamber_index.py

@@ -0,0 +1,304 @@
+"""
+PyTorch-compatible chamber lookup for H4 ChamberTree.
+
+Provides a bridge between PyTorch tensors (gradient-tracked) and the
+numpy-based H4ChamberTree (discrete, non-differentiable). The key trick:
+
+  - ChamberTree does fast O(log t) filtering to find top-k candidate keys
+  - We return candidate indices back to PyTorch
+  - Attention scores are computed only over candidates (differentiable)
+  - Gradients flow through Q/K projections and scores, not through the tree
+
+This gives O(k) attention per query where k << t.
+
+If the compiled Rust backend (h4_rust) is available, RustChamberIndex provides
+a much faster implementation. Falls back to pure-Python ChamberIndex otherwise.
+"""
+
+import numpy as np
+import torch
+from typing import List, Tuple, Optional
+import sys
+import os
+
+# Rust backend detection — optional, graceful fallback to Python
+try:
+    import h4_rust
+    RUST_AVAILABLE = True
+except ImportError:
+    RUST_AVAILABLE = False
+
+# Add parent to path for imports
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from h4_polytopic_attention import H4ChamberTree, build_coxeter_chambers, generate_600_cell_vertices
+
+
+class ChamberIndex:
+    """
+    Manages a set of H4ChamberTrees (one per head) and provides
+    batch top-k candidate lookup compatible with PyTorch autograd.
+    """
+
+    def __init__(self, n_heads: int, simple_roots: np.ndarray):
+        self.n_heads = n_heads
+        self.simple_roots = simple_roots
+        self.trees = [H4ChamberTree(simple_roots) for _ in range(n_heads)]
+        self._keys_by_head = [[] for _ in range(n_heads)]  # track insertion order
+
+    def reset(self):
+        """Clear all trees and rebuild."""
+        self.trees = [H4ChamberTree(self.simple_roots) for _ in range(self.n_heads)]
+        self._keys_by_head = [[] for _ in range(self.n_heads)]
+
+    def insert_keys(self, keys: torch.Tensor):
+        """
+        Insert keys for all heads at current timestep.
+
+        Args:
+            keys: (n_heads, 4) tensor of key vectors to insert
+        """
+        keys_np = keys.detach().cpu().numpy()
+        t = len(self._keys_by_head[0])  # current position index
+        for h in range(self.n_heads):
+            key = keys_np[h]
+            # Use position index as both value and timestamp
+            self.trees[h].insert(key, np.array([t], dtype=np.float64), t)
+            self._keys_by_head[h].append(key.copy())
+
+    def bulk_insert(self, keys: torch.Tensor):
+        """
+        Insert a full sequence of keys for all heads.
+
+        Args:
+            keys: (seq_len, n_heads, 4) tensor of key vectors
+        """
+        seq_len = keys.shape[0]
+        keys_np = keys.detach().cpu().numpy()
+        for t in range(seq_len):
+            for h in range(self.n_heads):
+                key = keys_np[t, h]
+                self.trees[h].insert(key, np.array([t], dtype=np.float64), t)
+                self._keys_by_head[h].append(key.copy())
+
+    def query_topk(
+        self,
+        queries: torch.Tensor,
+        k: int,
+        causal_mask_pos: Optional[int] = None,
+    ) -> List[List[List[int]]]:
+        """
+        For each query, find top-k candidate key indices using ChamberTree.
+
+        Args:
+            queries: (n_queries, n_heads, 4) tensor of query vectors
+            k: number of candidates per query per head
+            causal_mask_pos: if set, only return candidates with index <= this value
+
+        Returns:
+            List of shape [n_queries][n_heads][<=k] containing key indices.
+            These indices can be used to gather from the full key/value tensors.
+        """
+        n_queries = queries.shape[0]
+        queries_np = queries.detach().cpu().numpy()
+        results = []
+
+        for q_idx in range(n_queries):
+            head_results = []
+            for h in range(self.n_heads):
+                query = queries_np[q_idx, h]
+                # Query tree for top candidates
+                # Request more than k since some may be filtered by causal mask
+                tree_results = self.trees[h].query_max_dot(query, k=k * 2)
+
+                indices = []
+                for score, value, timestamp in tree_results:
+                    t_idx = int(value[0]) if len(value) > 0 else timestamp
+                    if causal_mask_pos is not None and t_idx > causal_mask_pos:
+                        continue
+                    indices.append(t_idx)
+                    if len(indices) >= k:
+                        break
+
+                # If tree didn't return enough, fall back to scanning
+                if len(indices) < k and len(self._keys_by_head[h]) > 0:
+                    max_pos = causal_mask_pos if causal_mask_pos is not None else len(self._keys_by_head[h]) - 1
+                    all_keys = np.array(self._keys_by_head[h][:max_pos + 1])
+                    if len(all_keys) > 0:
+                        dots = all_keys @ query
+                        sorted_idx = np.argsort(-dots)
+                        existing = set(indices)
+                        for idx in sorted_idx:
+                            if idx not in existing:
+                                indices.append(int(idx))
+                                existing.add(int(idx))
+                            if len(indices) >= k:
+                                break
+
+                head_results.append(indices)
+            results.append(head_results)
+
+        return results
+
+
+def compute_chamber_ids(keys: torch.Tensor, simple_roots: torch.Tensor) -> torch.Tensor:
+    """
+    Compute chamber IDs for a batch of keys (differentiable w.r.t. nothing,
+    but useful for logging chamber utilization).
+
+    Args:
+        keys: (..., 4) tensor of key vectors
+        simple_roots: (4, 4) tensor of H4 simple roots
+
+    Returns:
+        (...,) tensor of integer chamber IDs (0-15 for 4-bit sign pattern)
+    """
+    # Dot products with all 4 roots: (..., 4)
+    dots = keys @ simple_roots.T
+    # Sign pattern → 4-bit chamber ID
+    signs = (dots >= 0).long()
+    ids = signs[..., 0] * 8 + signs[..., 1] * 4 + signs[..., 2] * 2 + signs[..., 3]
+    return ids
+
+
+def chamber_utilization(chamber_ids: torch.Tensor, n_chambers: int = 16) -> dict:
+    """
+    Compute chamber utilization statistics.
+
+    Returns:
+        Dict with 'counts' (per-chamber), 'entropy' (Shannon entropy),
+        and 'max_ratio' (max/mean ratio, 1.0 = perfectly uniform).
+    """
+    counts = torch.zeros(n_chambers, dtype=torch.long, device=chamber_ids.device)
+    flat = chamber_ids.flatten()
+    for i in range(n_chambers):
+        counts[i] = (flat == i).sum()
+
+    total = counts.sum().float()
+    if total == 0:
+        return {'counts': counts, 'entropy': 0.0, 'max_ratio': 0.0}
+
+    probs = counts.float() / total
+    # Shannon entropy (nats)
+    log_probs = torch.where(probs > 0, torch.log(probs), torch.zeros_like(probs))
+    entropy = -(probs * log_probs).sum().item()
+
+    mean_count = total / n_chambers
+    max_ratio = (counts.max().float() / mean_count).item() if mean_count > 0 else 0.0
+
+    return {
+        'counts': counts,
+        'entropy': entropy,
+        'max_ratio': max_ratio,
+    }
+
+
+class RustChamberIndex:
+    """
+    Rust-accelerated chamber index using h4_rust compiled backend.
+    API-compatible with ChamberIndex for drop-in replacement.
+
+    All heavy computation (dot products, sorting, chamber indexing) runs
+    in compiled Rust via PyO3/numpy, typically 10-100x faster than Python.
+    """
+
+    def __init__(self, n_heads: int, simple_roots: np.ndarray):
+        if not RUST_AVAILABLE:
+            raise ImportError("h4_rust is not available. Install with: cd rust && maturin develop --release")
+        self.n_heads = n_heads
+        self.simple_roots = simple_roots  # (4, 4) numpy array
+        self._keys_by_head = [[] for _ in range(n_heads)]  # list of (4,) arrays per head
+
+    def reset(self):
+        """Clear all stored keys."""
+        self._keys_by_head = [[] for _ in range(self.n_heads)]
+
+    def insert_keys(self, keys: torch.Tensor):
+        """
+        Insert keys for all heads at current timestep.
+
+        Args:
+            keys: (n_heads, 4) tensor of key vectors to insert
+        """
+        keys_np = keys.detach().cpu().numpy()
+        for h in range(self.n_heads):
+            self._keys_by_head[h].append(keys_np[h].copy())
+
+    def bulk_insert(self, keys: torch.Tensor):
+        """
+        Insert a full sequence of keys for all heads.
+
+        Args:
+            keys: (seq_len, n_heads, 4) tensor of key vectors
+        """
+        keys_np = keys.detach().cpu().numpy()
+        seq_len = keys_np.shape[0]
+        for t in range(seq_len):
+            for h in range(self.n_heads):
+                self._keys_by_head[h].append(keys_np[t, h].copy())
+
+    def query_topk(
+        self,
+        queries: torch.Tensor,
+        k: int,
+        causal_mask_pos: Optional[int] = None,
+    ) -> List[List[List[int]]]:
+        """
+        For each query, find top-k candidate key indices using Rust backend.
+
+        Args:
+            queries: (n_queries, n_heads, 4) tensor of query vectors
+            k: number of candidates per query per head
+            causal_mask_pos: if set, only consider keys with index <= this value
+
+        Returns:
+            List of shape [n_queries][n_heads][<=k] containing key indices.
+        """
+        n_queries = queries.shape[0]
+        queries_np = queries.detach().cpu().numpy()
+        results = []
+
+        for q_idx in range(n_queries):
+            head_results = []
+            for h in range(self.n_heads):
+                n_keys = len(self._keys_by_head[h])
+                if n_keys == 0:
+                    head_results.append([])
+                    continue
+
+                # Apply causal mask: only use keys up to causal_mask_pos
+                max_pos = causal_mask_pos if causal_mask_pos is not None else n_keys - 1
+                effective_n = min(n_keys, max_pos + 1)
+
+                if effective_n == 0:
+                    head_results.append([])
+                    continue
+
+                keys_arr = np.array(self._keys_by_head[h][:effective_n], dtype=np.float64)
+                query_arr = queries_np[q_idx, h:h+1].astype(np.float64)
+
+                actual_k = min(k, effective_n)
+                indices = h4_rust.query_topk(keys_arr, query_arr, actual_k)
+                # indices is (1, actual_k), extract the list and filter -1s
+                idx_list = [int(i) for i in indices[0] if i >= 0]
+                head_results.append(idx_list)
+
+            results.append(head_results)
+
+        return results
+
+
+def get_chamber_index(n_heads: int, simple_roots: np.ndarray, prefer_rust: bool = True):
+    """
+    Factory function: returns RustChamberIndex if available, else ChamberIndex.
+
+    Args:
+        n_heads: number of attention heads
+        simple_roots: (4, 4) numpy array of H4 simple roots
+        prefer_rust: if True (default), use Rust backend when available
+
+    Returns:
+        ChamberIndex or RustChamberIndex instance
+    """
+    if prefer_rust and RUST_AVAILABLE:
+        return RustChamberIndex(n_heads, simple_roots)
+    return ChamberIndex(n_heads, simple_roots)