feat: implement lazy-loading for attention matrices

- Add MatrixCache class with 60-min TTL for storing attention/QKV matrices
- Modify response builder to cache matrices instead of including in payload
- Add new endpoint /analyze/research/attention/matrix for on-demand retrieval
- Include requestId in responses for cache lookup
- Reduces initial response from 400MB+ to ~500KB-1MB

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

backend/model_service.py

@@ -18,6 +18,9 @@ import numpy as np
 import logging
 from datetime import datetime
 import traceback
+import uuid
+from threading import Lock
+from time import time as time_now
 from .auth import verify_api_key
 from .instrumentation import ModelInstrumentor, InstrumentationData, TokenMetadata
 from .storage import ZarrStorage, generate_run_id
@@ -56,6 +59,63 @@ def sanitize_for_json(obj):
     else:
         return obj
 
+
+# Matrix cache for lazy loading (60 min TTL)
+class MatrixCache:
+    """
+    Thread-safe in-memory cache for attention matrices.
+    Stores Q/K/V and attention weights per (request_id, step, layer, head).
+    """
+    def __init__(self, ttl_seconds: int = 3600):
+        self._cache: Dict[str, Dict] = {}
+        self._timestamps: Dict[str, float] = {}
+        self._lock = Lock()
+        self._ttl = ttl_seconds
+
+    def store(self, request_id: str, step: int, layer: int, head: int, data: dict):
+        """Store matrix data for a specific head."""
+        key = f"{request_id}:{step}:{layer}:{head}"
+        with self._lock:
+            self._cache[key] = data
+            self._timestamps[key] = time_now()
+
+    def get(self, request_id: str, step: int, layer: int, head: int) -> Optional[dict]:
+        """Retrieve matrix data, returning None if expired or not found."""
+        key = f"{request_id}:{step}:{layer}:{head}"
+        with self._lock:
+            if key in self._cache:
+                if time_now() - self._timestamps[key] < self._ttl:
+                    return self._cache[key]
+                else:
+                    # Expired - clean up
+                    del self._cache[key]
+                    del self._timestamps[key]
+        return None
+
+    def cleanup_expired(self):
+        """Remove all expired entries from cache."""
+        with self._lock:
+            now = time_now()
+            expired = [k for k, t in self._timestamps.items() if now - t >= self._ttl]
+            for k in expired:
+                del self._cache[k]
+                del self._timestamps[k]
+            if expired:
+                logger.info(f"MatrixCache: cleaned up {len(expired)} expired entries")
+
+    def get_stats(self) -> dict:
+        """Return cache statistics."""
+        with self._lock:
+            return {
+                "entries": len(self._cache),
+                "ttl_seconds": self._ttl
+            }
+
+
+# Global matrix cache instance
+matrix_cache = MatrixCache(ttl_seconds=3600)  # 60 min TTL
+
+
 app = FastAPI(title="Visualisable.ai Model Service", version="0.1.0")
 
 # CORS configuration for local development and production
@@ -1507,6 +1567,9 @@ async def analyze_research_attention(request: Dict[str, Any], authenticated: boo
         import time
         start_time = time.time()
 
+        # Generate unique request ID for matrix cache lookup
+        request_id = str(uuid.uuid4())
+
         # Get parameters
         prompt = request.get("prompt", "def quicksort(arr):")
         max_tokens = request.get("max_tokens", 8)
@@ -1777,15 +1840,21 @@ async def analyze_research_attention(request: Dict[str, Any], authenticated: boo
                             k_matrix = qkv_captures[layer_idx]['k'][:, head_idx, :].float().numpy().tolist()
                             v_matrix = qkv_captures[layer_idx]['v'][:, head_idx, :].float().numpy().tolist()
 
+                        # Store matrices in cache for lazy loading (reduces response size)
+                        matrix_cache.store(request_id, step, layer_idx, head_idx, {
+                            "attention_weights": attention_matrix,
+                            "q_matrix": q_matrix,
+                            "k_matrix": k_matrix,
+                            "v_matrix": v_matrix
+                        })
+
+                        # Return only metadata (matrices fetched on-demand via /matrix endpoint)
                         critical_heads.append({
                             "head_idx": head_idx,
                             "entropy": entropy,
                             "avg_entropy": avg_entropy,  # Averaged over all query positions
                             "max_weight": max_weight,
-                            "attention_weights": attention_matrix,  # Full attention matrix for spreadsheet
-                            "q_matrix": q_matrix,  # [seq_len, head_dim]
-                            "k_matrix": k_matrix,
-                            "v_matrix": v_matrix,
+                            "has_matrices": attention_matrix is not None,  # Flag for frontend
                             "pattern": {
                                 "type": pattern_type,
                                 "confidence": confidence
@@ -1915,6 +1984,7 @@ async def analyze_research_attention(request: Dict[str, Any], authenticated: boo
 
         # Build response
         response = {
+            "requestId": request_id,  # For lazy-loading matrices via /matrix endpoint
             "prompt": prompt,
             "promptTokens": build_token_data(prompt_token_ids, prompt_tokens, "prompt"),
             "generatedTokens": build_token_data(generated_token_ids, generated_tokens, "generated"),
@@ -1922,7 +1992,7 @@ async def analyze_research_attention(request: Dict[str, Any], authenticated: boo
             "tokenAlternatives": token_alternatives_by_step,  # Top-k alternatives for each token
             "layersDataByStep": layer_data_by_token,  # Layer data for ALL generation steps
             "layersData": layer_data_by_token[-1] if layer_data_by_token else [],  # Keep for backward compatibility
-            "qkvData": qkv_by_layer_head,
+            "qkvData": {},  # Deprecated: matrices now lazy-loaded via /matrix endpoint
             "modelInfo": {
                 "numLayers": n_layers,
                 "numHeads": n_heads,
@@ -1969,12 +2039,15 @@ async def analyze_research_attention_stream(request: Dict[str, Any], authenticat
             import time
             start_time = time.time()
 
+            # Generate unique request ID for matrix cache lookup
+            request_id = str(uuid.uuid4())
+
             # Get parameters
             prompt = request.get("prompt", "def quicksort(arr):")
             max_tokens = request.get("max_tokens", 8)
             temperature = request.get("temperature", 0.7)
 
-            logger.info(f"[SSE] Research attention analysis: prompt_len={len(prompt)}, max_tokens={max_tokens}")
+            logger.info(f"[SSE] Research attention analysis: prompt_len={len(prompt)}, max_tokens={max_tokens}, request_id={request_id}")
 
             # === STAGE 1: TOKENIZING ===
             yield sse_event('tokenizing', stage=1, totalStages=5, progress=2,
@@ -2233,15 +2306,21 @@ async def analyze_research_attention_stream(request: Dict[str, Any], authenticat
                                 k_matrix = qkv_captures[layer_idx]['k'][:, head_idx, :].float().numpy().tolist()
                                 v_matrix = qkv_captures[layer_idx]['v'][:, head_idx, :].float().numpy().tolist()
 
+                            # Store matrices in cache for lazy loading (reduces response size)
+                            matrix_cache.store(request_id, step, layer_idx, head_idx, {
+                                "attention_weights": attention_matrix,
+                                "q_matrix": q_matrix,
+                                "k_matrix": k_matrix,
+                                "v_matrix": v_matrix
+                            })
+
+                            # Return only metadata (matrices fetched on-demand via /matrix endpoint)
                             critical_heads.append({
                                 "head_idx": head_idx,
                                 "entropy": entropy,
                                 "avg_entropy": avg_entropy,  # Averaged over all query positions
                                 "max_weight": max_weight,
-                                "attention_weights": attention_matrix,
-                                "q_matrix": q_matrix,
-                                "k_matrix": k_matrix,
-                                "v_matrix": v_matrix,
+                                "has_matrices": attention_matrix is not None,  # Flag for frontend
                                 "pattern": {"type": pattern_type, "confidence": confidence} if pattern_type else None
                             })
 
@@ -2364,6 +2443,7 @@ async def analyze_research_attention_stream(request: Dict[str, Any], authenticat
 
             # Build response
             response = {
+                "requestId": request_id,  # For lazy-loading matrices via /matrix endpoint
                 "prompt": prompt,
                 "promptTokens": build_token_data(prompt_token_ids, prompt_tokens, "prompt"),
                 "generatedTokens": build_token_data(generated_token_ids, generated_tokens, "generated"),
@@ -2371,7 +2451,7 @@ async def analyze_research_attention_stream(request: Dict[str, Any], authenticat
                 "tokenAlternatives": token_alternatives_by_step,
                 "layersDataByStep": layer_data_by_token,
                 "layersData": layer_data_by_token[-1] if layer_data_by_token else [],
-                "qkvData": qkv_by_layer_head,
+                "qkvData": {},  # Deprecated: matrices now lazy-loaded via /matrix endpoint
                 "modelInfo": {
                     "numLayers": n_layers,
                     "numHeads": n_heads,
@@ -2418,6 +2498,50 @@ async def analyze_research_attention_stream(request: Dict[str, Any], authenticat
     )
 
 
+@app.get("/analyze/research/attention/matrix")
+async def get_attention_matrix(
+    request_id: str,
+    step: int,
+    layer: int,
+    head: int,
+    authenticated: bool = Depends(verify_api_key)
+):
+    """
+    Retrieve cached attention/QKV matrices for a specific head.
+
+    Used for lazy-loading matrix data when user clicks "View Matrix" in the frontend.
+    Matrices are cached during the initial analysis and available for 60 minutes.
+
+    Parameters:
+    - request_id: UUID from the original analysis response
+    - step: Generation step (0 = first generated token)
+    - layer: Layer index (0-based)
+    - head: Head index (0-based)
+
+    Returns:
+    - attention_weights: [seq_len, seq_len] attention matrix
+    - q_matrix: [seq_len, head_dim] query projections
+    - k_matrix: [seq_len, head_dim] key projections
+    - v_matrix: [seq_len, head_dim] value projections
+    """
+    data = matrix_cache.get(request_id, step, layer, head)
+    if data is None:
+        logger.warning(f"Matrix cache miss: request_id={request_id}, step={step}, layer={layer}, head={head}")
+        raise HTTPException(
+            status_code=404,
+            detail="Matrix data not found. Cache may have expired (60 min TTL). Please re-analyze."
+        )
+
+    logger.info(f"Matrix cache hit: request_id={request_id}, step={step}, layer={layer}, head={head}")
+    return data
+
+
+@app.get("/analyze/research/attention/matrix/stats")
+async def get_matrix_cache_stats(authenticated: bool = Depends(verify_api_key)):
+    """Return matrix cache statistics for monitoring."""
+    return matrix_cache.get_stats()
+
+
 @app.post("/analyze/study")
 async def analyze_study(request: StudyRequest, authenticated: bool = Depends(verify_api_key)):
     """