Merge: F3a — derive aggregate-row num_heads from request metadata (#1)

backend/model_service.py

@@ -71,6 +71,12 @@ class MatrixCache:
         self._cache: Dict[str, Dict] = {}
         self._timestamps: Dict[str, float] = {}
         self._request_ids: set = set()  # Track active request IDs
+        # Per-request metadata captured at analysis time. Aggregation across
+        # heads (e.g. mean/max in /attention/row) needs num_heads, and that
+        # number belongs to the model that produced the cached attention,
+        # not whichever model is loaded when the row is later retrieved.
+        # Storing it here decouples retrieval from `manager.model` state.
+        self._request_meta: Dict[str, Dict] = {}
         self._lock = Lock()
         self._ttl = ttl_seconds
 
@@ -83,6 +89,7 @@ class MatrixCache:
                 if k in self._timestamps:
                     del self._timestamps[k]
             self._request_ids.discard(request_id)
+            self._request_meta.pop(request_id, None)
             if keys_to_delete:
                 logger.info(f"MatrixCache: cleared {len(keys_to_delete)} entries for request {request_id[:8]}")
 
@@ -99,6 +106,10 @@ class MatrixCache:
                         del self._timestamps[k]
                 total_cleared += len(keys_to_delete)
             self._request_ids = {keep_request_id} if keep_request_id else set()
+            # Drop metadata for cleared requests; keep only the surviving one.
+            for rid in list(self._request_meta.keys()):
+                if rid != keep_request_id:
+                    self._request_meta.pop(rid, None)
             if total_cleared:
                 logger.info(f"MatrixCache: cleared {total_cleared} entries from old requests")
                 # Force garbage collection to release memory back to system
@@ -118,6 +129,29 @@ class MatrixCache:
             self._timestamps[key] = time_now()
             self._request_ids.add(request_id)
 
+    def set_request_metadata(self, request_id: str, num_heads: int,
+                             num_layers: Optional[int] = None,
+                             model_id: Optional[str] = None):
+        """
+        Record metadata for the model that produced this request's cached
+        matrices. Aggregation operations (e.g. mean/max across heads) read
+        num_heads from here rather than from the live `manager.model`, so
+        that switching models mid-session cannot corrupt aggregate results
+        for older requests.
+        """
+        with self._lock:
+            self._request_meta[request_id] = {
+                "num_heads": num_heads,
+                "num_layers": num_layers,
+                "model_id": model_id,
+            }
+
+    def get_request_metadata(self, request_id: str) -> Optional[dict]:
+        """Return per-request metadata recorded at analysis time, or None."""
+        with self._lock:
+            meta = self._request_meta.get(request_id)
+            return dict(meta) if meta else None
+
     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}"
@@ -169,20 +203,33 @@ class MatrixCache:
         return list(last_row)
 
     def get_aggregate_row(self, request_id: str, step: int, layer: int,
-                          num_heads: int, mode: str = "mean") -> Optional[list]:
+                          mode: str = "mean",
+                          num_heads: Optional[int] = None) -> Optional[list]:
         """
         Compute aggregated attention row across all heads for a layer.
 
+        num_heads is derived from per-request metadata recorded at analysis
+        time so the aggregation always uses the head count of the model that
+        produced the cached attention. The explicit `num_heads` argument is
+        retained for callers that want to override (e.g. tests); when None,
+        falls back to the request's stored metadata.
+
         Args:
             request_id: UUID from analysis
             step: Generation step
             layer: Layer index
-            num_heads: Number of attention heads in model
             mode: Aggregation mode - "mean" or "max"
+            num_heads: Optional override; otherwise derived from request meta
 
         Returns:
             List of aggregated attention weights, or None if data unavailable
+            or request metadata is missing.
         """
+        if num_heads is None:
+            meta = self.get_request_metadata(request_id)
+            if meta is None or meta.get("num_heads") is None:
+                return None
+            num_heads = meta["num_heads"]
         rows = []
         for h in range(num_heads):
             row = self.get_attention_row(request_id, step, layer, h)
@@ -1799,6 +1846,17 @@ async def analyze_research_attention(request: Dict[str, Any], authenticated: boo
         # Clear old cached matrices to free memory before starting new analysis
         matrix_cache.clear_old_requests(request_id)
 
+        # Record per-request model metadata so aggregate-row lookups can
+        # derive num_heads from this request rather than from the live
+        # `manager.model` (which may have been switched in the meantime).
+        if manager.adapter is not None:
+            matrix_cache.set_request_metadata(
+                request_id,
+                num_heads=manager.adapter.get_num_heads(),
+                num_layers=manager.adapter.get_num_layers(),
+                model_id=manager.model_id,
+            )
+
         # Get parameters
         prompt = request.get("prompt", "def quicksort(arr):")
         max_tokens = request.get("max_tokens", 8)
@@ -2586,6 +2644,17 @@ async def analyze_research_attention_stream(request: Dict[str, Any], authenticat
             # Clear old cached matrices to free memory before starting new analysis
             matrix_cache.clear_old_requests(request_id)
 
+            # Record per-request model metadata so aggregate-row lookups can
+            # derive num_heads from this request rather than from the live
+            # `manager.model` (which may have been switched in the meantime).
+            if manager.adapter is not None:
+                matrix_cache.set_request_metadata(
+                    request_id,
+                    num_heads=manager.adapter.get_num_heads(),
+                    num_layers=manager.adapter.get_num_layers(),
+                    model_id=manager.model_id,
+                )
+
             # Get parameters
             prompt = request.get("prompt", "def quicksort(arr):")
             max_tokens = request.get("max_tokens", 8)
@@ -3845,15 +3914,13 @@ async def get_attention_row(
     - head: Head index (null if aggregated)
     - aggregate_mode: Mode used if aggregated (null otherwise)
     """
-    # Get number of heads from model config
-    if not manager.model:
-        raise HTTPException(status_code=503, detail="Model not loaded")
-
-    config = manager.model.config
-    num_heads = getattr(config, 'num_attention_heads', getattr(config, 'n_head', 16))
-
+    # Aggregation parameters (num_heads in particular) come from per-request
+    # metadata recorded at analysis time, not from the live model config.
+    # This keeps aggregate rows correct even after the user switches models
+    # mid-session — the row is averaged over the heads of the model that
+    # actually produced the cached attention.
     if head is not None:
-        # Fetch specific head
+        # Fetch specific head — no aggregation, no metadata lookup needed.
         attention_row = matrix_cache.get_attention_row(request_id, step, layer, head)
         if attention_row is None:
             logger.warning(f"Attention row cache miss: request_id={request_id}, step={step}, layer={layer}, head={head}")
@@ -3870,9 +3937,10 @@ async def get_attention_row(
             "aggregate_mode": None
         }
     else:
-        # Aggregate across all heads
+        # Aggregate across all heads. num_heads is derived inside
+        # get_aggregate_row from the request's stored metadata.
         attention_row = matrix_cache.get_aggregate_row(
-            request_id, step, layer, num_heads, aggregate_mode
+            request_id, step, layer, aggregate_mode
         )
         if attention_row is None:
             logger.warning(f"Attention row aggregate cache miss: request_id={request_id}, step={step}, layer={layer}")