Emit per-token embedding L2 norms for RQ1 Step 3

Adds a `_compute_embedding_norms` helper that bounds memory by
unique token IDs (`torch.unique` + `weight.index_select`) so long
contexts don't allocate `[1, seq_len, hidden_size]` × 2 unnecessarily.

Both analyse endpoints now return `embeddingNorms: List[float]`
covering the prompt + generated tokens, used by the new code-first
Mechanism lens to surface real per-token embedding magnitudes
instead of synthetic illustrative values.

Also extends MatrixCache with `get_aggregate_attention_matrix()`
(mean / max pooling across heads) so the Step 4 attention view can
fetch the aggregate matrix from a single endpoint instead of
synthesising it on the client.

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>

backend/model_service.py

@@ -283,6 +283,56 @@ class MatrixCache:
             # Default to mean for unknown modes
             return np.mean(arr, axis=0).tolist()
 
+    def get_aggregate_attention_matrix(
+        self,
+        request_id: str,
+        step: int,
+        layer: int,
+        mode: str = "mean",
+        num_heads: Optional[int] = None,
+    ) -> Optional[list]:
+        """
+        Compute the aggregated attention MATRIX across all heads for a
+        layer. Mirrors get_aggregate_row but returns the full
+        [seq_len, seq_len] matrix instead of one row.
+
+        Used by the Mechanism lens's Step 4 (Attention) to render the
+        per-step matrix in its "all heads (mean)" default mode without
+        forcing the frontend to fetch + aggregate N matrices itself.
+
+        Args:
+            request_id: UUID from analysis
+            step: Generation step
+            layer: Layer index
+            mode: Aggregation mode - "mean" or "max"
+            num_heads: Optional override; otherwise derived from request meta
+
+        Returns:
+            Aggregated 2-D matrix as a list of lists, 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"]
+        matrices = []
+        for h in range(num_heads):
+            data = self.get(request_id, step, layer, h)
+            if data is None:
+                continue
+            attn = data.get("attention_weights")
+            if attn is None:
+                continue
+            matrices.append(np.asarray(attn))
+        if not matrices:
+            return None
+        arr = np.stack(matrices, axis=0)  # (n_heads, seq, seq)
+        if mode == "max":
+            return np.max(arr, axis=0).tolist()
+        # Default to mean for "mean" or unknown modes.
+        return np.mean(arr, axis=0).tolist()
+
 
 # Global matrix cache instance
 matrix_cache = MatrixCache(ttl_seconds=3600)  # 60 min TTL
@@ -300,6 +350,52 @@ def _classify_stability(margin: float) -> str:
         return "fragile"
 
 
+def _compute_embedding_norms(manager, token_ids: List[int]) -> List[float]:
+    """
+    Per-token L2 norms of the layer-0 embedding vectors, indexed by the
+    token ids supplied. Used by the analyse endpoints to surface a
+    real-data anchor for the residual-norm trace shown in the
+    Mechanism lens's Microscope rail (RQ1).
+
+    Memory is bounded by the number of *unique* token ids in the
+    sequence rather than the sequence length: a long prompt with
+    repeated tokens reuses the same embedding row, so we gather rows
+    once for unique ids and scatter the resulting scalar back to the
+    sequence positions. This avoids materialising a full
+    [1, seq_len, hidden_size] activation tensor — important for the
+    long-context models this service supports (e.g. devstral-small at
+    131,072 tokens, where the naive path would allocate multi-GB at
+    end-of-request and risk OOM).
+
+    Failures (missing module, unexpected layout) downgrade gracefully
+    to an empty list so the rest of the trace is unaffected.
+    """
+    if not token_ids:
+        return []
+    try:
+        with torch.no_grad():
+            embed_layer = manager.model.get_input_embeddings()
+            ids_tensor = torch.tensor(
+                token_ids, dtype=torch.long, device=manager.device
+            )
+            unique_ids, inverse = torch.unique(
+                ids_tensor, return_inverse=True
+            )
+            # Index the embedding weight matrix directly. Allocation is
+            # (n_unique, hidden_size) instead of (seq_len, hidden_size)
+            # — typically a few hundred rows for code-generation
+            # prompts regardless of how long the prompt is.
+            unique_rows = embed_layer.weight.index_select(0, unique_ids)
+            unique_norms = torch.linalg.vector_norm(
+                unique_rows.float(), dim=-1
+            )
+            seq_norms = unique_norms[inverse]
+            return [float(v) for v in seq_norms.cpu().tolist()]
+    except Exception as e:  # pragma: no cover — defensive
+        logger.warning(f"Failed to compute embedding norms: {e}")
+        return []
+
+
 class HiddenStateCache:
     """
     Cache for hidden states and logits per (request_id, step).
@@ -2654,6 +2750,16 @@ async def analyze_research_attention(request: Dict[str, Any], authenticated: boo
                 })
             return result
 
+        # Embedding L2 norms — see streaming endpoint for the full
+        # rationale (RQ1 layer-0 anchor for the residual-norm trace).
+        # Covers both prompt tokens and generated tokens so the
+        # Mechanism lens can render bars for the full context window
+        # (prompt + previously-generated tokens) at any step.
+        embedding_norms: List[float] = _compute_embedding_norms(
+            manager,
+            list(prompt_token_ids) + list(generated_token_ids),
+        )
+
         # Build response
         response = {
             "requestId": request_id,  # For lazy-loading matrices via /matrix endpoint
@@ -2672,6 +2778,7 @@ async def analyze_research_attention(request: Dict[str, Any], authenticated: boo
                 "headDim": head_dim,
                 "vocabSize": manager.model.config.vocab_size
             },
+            "embeddingNorms": embedding_norms,
             "generationTime": generation_time,
             "numTokensGenerated": len(generated_tokens)
         }
@@ -3953,6 +4060,22 @@ async def analyze_research_attention_stream(request: Dict[str, Any], authenticat
                     "flip_count": tuned_flip_count,
                 }
 
+            # === Embedding L2 norms — per-input-token layer-0 anchor ===
+            # One float per token in the FULL context (prompt +
+            # generated): ‖e_i‖ where e_i is the row the embedding
+            # matrix returns for token i. The Mechanism lens uses this
+            # as the real-data anchor for the residual-norm trace
+            # shown in the Microscope rail (RQ1: developer-
+            # interpretable architectural signal at the embedding
+            # stage). Generated tokens are included so the chart
+            # reflects the actual input the model sees at every
+            # selected step. Memory is bounded by the number of
+            # *unique* token ids — see _compute_embedding_norms.
+            embedding_norms: List[float] = _compute_embedding_norms(
+                manager,
+                list(prompt_token_ids) + list(generated_token_ids),
+            )
+
             # Build response
             response = {
                 "requestId": request_id,  # For lazy-loading matrices via /matrix endpoint
@@ -3974,6 +4097,7 @@ async def analyze_research_attention_stream(request: Dict[str, Any], authenticat
                     "ffnType": ffn_type,
                     "intermediateSize": getattr(manager.model.config, 'intermediate_size', None),
                 },
+                "embeddingNorms": embedding_norms,
                 "generationTime": generation_time,
                 "numTokensGenerated": len(generated_tokens),
                 "marginStats": margin_stats,
@@ -4021,46 +4145,87 @@ async def get_attention_matrix(
     request_id: str,
     step: int,
     layer: int,
-    head: int,
-    authenticated: bool = Depends(verify_api_key)
+    head: Optional[int] = None,
+    aggregate_mode: str = "mean",
+    authenticated: bool = Depends(verify_api_key),
 ):
     """
-    Retrieve cached attention/QKV matrices for a specific head.
+    Retrieve cached attention/QKV matrices for a specific head, OR an
+    aggregated attention matrix across all heads when `head` is omitted.
 
-    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.
+    Used for lazy-loading matrix data when the user opens an inline
+    head panel (per-head Q/K/V) or the Mechanism lens's Step 4
+    attention matrix view (aggregate by default). 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)
+    - head: Head index (0-based). Omit for an aggregate matrix across
+      all heads — the returned payload then contains only
+      `attention_weights` (the Q/K/V projections are per-head and
+      have no canonical aggregate form).
+    - aggregate_mode: "mean" or "max" when `head` is omitted.
 
     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
+    - q_matrix / k_matrix / v_matrix: [seq_len, head_dim] projections
+      (only when `head` is specified)
+    - layer: Layer index
+    - head: Head index, or null when aggregated
+    - aggregate_mode: Aggregation mode used, or null for per-head
     """
-    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}")
+    if head is not None:
+        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}"
+        )
+        # Convert numpy arrays to lists for JSON serialization. Arrays
+        # are stored as numpy for memory efficiency, converted on-
+        # demand here.
+        response = {}
+        for key, value in data.items():
+            if value is not None and hasattr(value, "tolist"):
+                response[key] = value.tolist()
+            else:
+                response[key] = value
+        response["layer"] = layer
+        response["head"] = head
+        response["aggregate_mode"] = None
+        return response
+
+    # Aggregate path — averages or max-pools attention matrices across
+    # heads. Q/K/V projections are intentionally omitted; they're
+    # per-head quantities and don't have a canonical aggregate.
+    aggregate = matrix_cache.get_aggregate_attention_matrix(
+        request_id, step, layer, aggregate_mode
+    )
+    if aggregate is None:
+        logger.warning(
+            f"Aggregate matrix cache miss: request_id={request_id}, step={step}, layer={layer}, mode={aggregate_mode}"
+        )
         raise HTTPException(
             status_code=404,
-            detail="Matrix data not found. Cache may have expired (60 min TTL). Please re-analyze."
+            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}")
-
-    # Convert numpy arrays to lists for JSON serialization
-    # Arrays are stored as numpy for memory efficiency, converted on-demand here
-    response = {}
-    for key, value in data.items():
-        if value is not None and hasattr(value, 'tolist'):
-            response[key] = value.tolist()
-        else:
-            response[key] = value
-    return response
+    logger.info(
+        f"Aggregate matrix cache hit: request_id={request_id}, step={step}, layer={layer}, mode={aggregate_mode}"
+    )
+    return {
+        "attention_weights": aggregate,
+        "layer": layer,
+        "head": None,
+        "aggregate_mode": aggregate_mode,
+    }
 
 
 @app.get("/analyze/research/attention/matrix/stats")