Merge pull request #6 from gary-boon/feature/rq1-ffn-intermediate

Backend: emit per-token embedding norms, top-K SwiGLU neurons, unembedding cosine + Devstral tokenizer fallback

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).
@@ -581,8 +677,24 @@ class ModelManager:
                 attn_implementation="eager"
             ).to(self.device)
             
-            # Load tokenizer
-            self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+            # Load tokenizer. Devstral's official repo ships only tekken.json,
+            # which AutoTokenizer can't consume on tokenizers ≥ 0.20. Fall back
+            # to the Unsloth mirror, which carries the same Tekken vocabulary
+            # in standard HF format. Same vocab size, same token IDs, same
+            # special tokens — only the on-disk file format differs.
+            try:
+                self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
+            except (OSError, EnvironmentError) as tok_err:
+                if self.model_id == "devstral-small":
+                    logger.warning(
+                        f"AutoTokenizer failed for {self.model_name} ({tok_err}); "
+                        f"falling back to unsloth/Devstral-Small-2507 (same Tekken vocab, HF format)."
+                    )
+                    self.tokenizer = AutoTokenizer.from_pretrained(
+                        "unsloth/Devstral-Small-2507"
+                    )
+                else:
+                    raise
             # Set pad_token if the tokenizer allows it (some like MistralCommonTokenizer don't)
             try:
                 self.tokenizer.pad_token = self.tokenizer.eos_token
@@ -2103,16 +2215,85 @@ async def analyze_research_attention(request: Dict[str, Any], authenticated: boo
                 top_n_display = 10  # Get top 10 alternatives for display
                 top_raw_logits, top_raw_indices = torch.topk(raw_logits, k=min(top_n_display, len(raw_logits)))
 
+                # ── Unembedding-cosine prep ──────────────────────────
+                # Each logit is W_E[v] · h_post_norm where h_post_norm
+                # is the post-final-norm residual at the last position.
+                # cos(W_E[v], h_post_norm) is the *direction-only*
+                # alignment — strips out magnitude so a developer can
+                # see whether a high-logit token won by alignment or
+                # by magnitude. Computed once per step (cheap; one
+                # final-norm call + d_model dot product per token).
+                cos_h_norm = None
+                W_E = None
+                try:
+                    final_hidden_pre_norm = outputs.hidden_states[-1][0, -1, :]
+                    if hasattr(manager.model, 'model') and hasattr(manager.model.model, 'norm'):
+                        h_post_norm = manager.model.model.norm(
+                            final_hidden_pre_norm.unsqueeze(0)
+                        )[0]
+                    elif hasattr(manager.model, 'transformer') and hasattr(manager.model.transformer, 'ln_f'):
+                        h_post_norm = manager.model.transformer.ln_f(
+                            final_hidden_pre_norm.unsqueeze(0)
+                        )[0]
+                    else:
+                        h_post_norm = final_hidden_pre_norm
+                    h_norm_val = float(torch.linalg.vector_norm(h_post_norm).item())
+                    if h_norm_val > 1e-9:
+                        cos_h_norm = h_post_norm / h_norm_val
+                        if hasattr(manager.model, 'lm_head'):
+                            W_E = manager.model.lm_head.weight
+                        else:
+                            logger.warning(
+                                "[cosine prep] manager.model has no lm_head — cosine emission skipped"
+                            )
+                    else:
+                        logger.warning(
+                            f"[cosine prep] h_post_norm magnitude too small: {h_norm_val}"
+                        )
+                    if cos_h_norm is not None and W_E is not None and step == 0:
+                        # One-time confirmation per generation that the
+                        # cosine path is live.
+                        logger.info(
+                            f"[cosine prep] enabled · h_post_norm shape {tuple(h_post_norm.shape)} · "
+                            f"W_E shape {tuple(W_E.shape)} · h_norm {h_norm_val:.3f}"
+                        )
+                except Exception as e:
+                    logger.warning(
+                        f"[cosine prep] failed at step {step}: {type(e).__name__}: {e}"
+                    )
+                    cos_h_norm = None
+                    W_E = None
+
+                def _cos_for(token_idx: int):
+                    """cosine(W_E[token_idx], h_post_norm) and ‖W_E row‖."""
+                    if cos_h_norm is None or W_E is None:
+                        return None, None
+                    try:
+                        row = W_E[token_idx]
+                        row_norm = float(torch.linalg.vector_norm(row).item())
+                        if row_norm < 1e-9:
+                            return None, row_norm
+                        cos = float(((row / row_norm) @ cos_h_norm).item())
+                        return cos, row_norm
+                    except Exception:
+                        return None, None
+
                 # Build raw logits entries (before temperature)
                 logits_entries = []
                 for rank, (logit_val, idx) in enumerate(zip(top_raw_logits.tolist(), top_raw_indices.tolist())):
                     token_text = manager.tokenizer.decode([idx], skip_special_tokens=False)
-                    logits_entries.append({
+                    cos, row_norm = _cos_for(idx)
+                    entry = {
                         "token": token_text,
                         "token_id": idx,
                         "logit": logit_val,
-                        "rank": rank + 1
-                    })
+                        "rank": rank + 1,
+                    }
+                    if cos is not None:
+                        entry["cosine_sim"] = round(cos, 4)
+                    if row_norm is not None:
+                        entry["unemb_row_norm"] = round(row_norm, 4)
+                    logits_entries.append(entry)
 
                 # Greedy token (argmax of raw logits, before any sampling)
                 greedy_token_id = torch.argmax(raw_logits).item()
@@ -2236,13 +2417,21 @@ async def analyze_research_attention(request: Dict[str, Any], authenticated: boo
                     })
                     # Also add to logits if not present
                     if next_token_id not in [e["token_id"] for e in logits_entries]:
-                        logits_entries.append({
+                        cos_outlier, row_norm_outlier = _cos_for(next_token_id)
+                        outlier_entry = {
                             "token": next_token_text,
                             "token_id": next_token_id,
                             "logit": selected_logit,
                             "rank": selected_rank,
-                            "is_selected_outlier": True
-                        })
+                            "is_selected_outlier": True,
+                        }
+                        if cos_outlier is not None:
+                            outlier_entry["cosine_sim"] = round(cos_outlier, 4)
+                        if row_norm_outlier is not None:
+                            outlier_entry["unemb_row_norm"] = round(
+                                row_norm_outlier, 4
+                            )
+                        logits_entries.append(outlier_entry)
 
                 # Build sampling metadata. `is_filtered` and `eligible_count`
                 # let the frontend decide whether to show one or two
@@ -2654,6 +2843,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 +2871,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)
         }
@@ -2921,6 +3121,16 @@ async def analyze_research_attention_stream(request: Dict[str, Any], authenticat
             attn_output_norms = {}
             mlp_output_norms = {}
             gate_activation_stats = {}
+            # Per-layer top-K SwiGLU intermediate activations at the
+            # predicting position (last token of the current forward
+            # pass). Each entry is a list of {idx, value} for the
+            # top-K neurons by |silu(W_g·x) ⊙ (W_u·x)| at that layer
+            # for that step. Powers RQ1 Step 5's "Behind the SwiGLU"
+            # per-neuron view without storing the full 14,336-d
+            # vector — bounded memory, full transparency for the
+            # most-loaded neurons.
+            ffn_top_neurons = {}
+            FFN_TOPK = 32
 
             def make_attn_output_hook(layer_idx):
                 def hook(module, input, output):
@@ -2965,6 +3175,63 @@ async def analyze_research_attention_stream(request: Dict[str, Any], authenticat
                         pass
                 return hook
 
+            def make_ffn_top_neurons_hook(layer_idx):
+                """
+                Capture the top-K of the SwiGLU intermediate activation
+                vector at the predicting (last) position for this layer.
+
+                Computes silu(W_gate · x) ⊙ (W_up · x) — the same
+                14,336-d intermediate that gets fed to W_down. Storing
+                only the top-K (=32) by |value| keeps memory bounded:
+                32 × {idx:int, value:float} ≈ 256 bytes per layer per
+                step, vs 14,336 × 4 bytes = 57KB per layer per step
+                if we stored the full vector. The top-K is what a
+                developer reads anyway (the rest are noise) and the
+                stored neuron index lets the frontend label "neuron
+                8421 dominated this token at L37".
+
+                Only fires for SwiGLU layers (gate_proj + up_proj
+                present); other architectures are skipped at
+                registration time.
+                """
+                def hook(module, input, output):
+                    try:
+                        inp = input[0] if isinstance(input, tuple) else input
+                        if inp.dim() == 3:
+                            inp = inp[0, -1]  # last (predicting) token
+                        elif inp.dim() == 2:
+                            inp = inp[-1]
+                        if hasattr(module, 'gate_proj') and hasattr(module, 'up_proj'):
+                            gate_out = torch.nn.functional.silu(
+                                module.gate_proj(inp)
+                            )
+                            up_out = module.up_proj(inp)
+                            intermediate = gate_out * up_out
+                            abs_inter = intermediate.abs()
+                            k = min(FFN_TOPK, intermediate.shape[-1])
+                            top_vals, top_idx = torch.topk(abs_inter, k=k)
+                            # Use the signed values (not abs) so the
+                            # frontend can render direction; abs only
+                            # drove the ranking.
+                            signed = intermediate[top_idx]
+                            ffn_top_neurons[layer_idx] = {
+                                "k": k,
+                                "intermediate_size": int(intermediate.shape[-1]),
+                                "neurons": [
+                                    {
+                                        "idx": int(i),
+                                        "value": round(float(v), 4),
+                                    }
+                                    for i, v in zip(
+                                        top_idx.cpu().tolist(),
+                                        signed.cpu().tolist(),
+                                    )
+                                ],
+                            }
+                    except Exception:
+                        pass
+                return hook
+
             # Cache for decoded token texts (reused across heads within a step)
             step_token_texts_cache: Dict[str, Any] = {}
 
@@ -2996,6 +3263,14 @@ async def analyze_research_attention_stream(request: Dict[str, Any], authenticat
                                 hooks.append(hook)
                                 if layer_idx == 0:
                                     ffn_type = "swiglu"
+                            # FFN top-K intermediate hook — needs both
+                            # gate_proj and up_proj to compute the
+                            # SwiGLU intermediate vector.
+                            if hasattr(layer.mlp, 'gate_proj') and hasattr(layer.mlp, 'up_proj'):
+                                hook = layer.mlp.register_forward_hook(
+                                    make_ffn_top_neurons_hook(layer_idx)
+                                )
+                                hooks.append(hook)
                 logger.info(f"Registered attn/MLP output hooks for contribution tracking (ffn_type={ffn_type})")
             except Exception as hook_error:
                 logger.warning(f"Could not register attn/MLP hooks: {hook_error}")
@@ -3017,6 +3292,7 @@ async def analyze_research_attention_stream(request: Dict[str, Any], authenticat
                     attn_output_norms.clear()
                     mlp_output_norms.clear()
                     gate_activation_stats.clear()
+                    ffn_top_neurons.clear()
 
                     # Forward pass with full outputs
                     outputs = manager.model(
@@ -3058,15 +3334,81 @@ async def analyze_research_attention_stream(request: Dict[str, Any], authenticat
                         else:
                             return manager.tokenizer.decode([tid], skip_special_tokens=False)
 
+                    # ── Unembedding-cosine prep (SSE path) ──────────
+                    # Mirrors the non-SSE path. Each logit is
+                    # W_E[v] · h_post_norm; cos(W_E[v], h_post_norm)
+                    # strips magnitude so a developer can read whether
+                    # a high-logit token won by direction or by
+                    # row magnitude. Cheap (one matmul per token per
+                    # step) and computed once per step.
+                    cos_h_norm = None
+                    W_E = None
+                    try:
+                        final_hidden_pre_norm = outputs.hidden_states[-1][0, -1, :]
+                        if hasattr(manager.model, 'model') and hasattr(manager.model.model, 'norm'):
+                            h_post_norm = manager.model.model.norm(
+                                final_hidden_pre_norm.unsqueeze(0)
+                            )[0]
+                        elif hasattr(manager.model, 'transformer') and hasattr(manager.model.transformer, 'ln_f'):
+                            h_post_norm = manager.model.transformer.ln_f(
+                                final_hidden_pre_norm.unsqueeze(0)
+                            )[0]
+                        else:
+                            h_post_norm = final_hidden_pre_norm
+                        h_norm_val = float(torch.linalg.vector_norm(h_post_norm).item())
+                        if h_norm_val > 1e-9:
+                            cos_h_norm = h_post_norm / h_norm_val
+                            if hasattr(manager.model, 'lm_head'):
+                                W_E = manager.model.lm_head.weight
+                            else:
+                                logger.warning(
+                                    "[SSE cosine prep] manager.model has no lm_head — cosine emission skipped"
+                                )
+                        else:
+                            logger.warning(
+                                f"[SSE cosine prep] h_post_norm magnitude too small: {h_norm_val}"
+                            )
+                        if cos_h_norm is not None and W_E is not None and step == 0:
+                            logger.info(
+                                f"[SSE cosine prep] enabled · h_post_norm shape {tuple(h_post_norm.shape)} · "
+                                f"W_E shape {tuple(W_E.shape)} · h_norm {h_norm_val:.3f}"
+                            )
+                    except Exception as e:
+                        logger.warning(
+                            f"[SSE cosine prep] failed at step {step}: {type(e).__name__}: {e}"
+                        )
+                        cos_h_norm = None
+                        W_E = None
+
+                    def _cos_for(token_idx: int):
+                        """cosine(W_E[token_idx], h_post_norm) and ‖W_E row‖."""
+                        if cos_h_norm is None or W_E is None:
+                            return None, None
+                        try:
+                            row = W_E[token_idx]
+                            row_norm = float(torch.linalg.vector_norm(row).item())
+                            if row_norm < 1e-9:
+                                return None, row_norm
+                            cos = float(((row / row_norm) @ cos_h_norm).item())
+                            return cos, row_norm
+                        except Exception:
+                            return None, None
+
                     logits_entries = []
                     for rank, (logit_val, idx) in enumerate(zip(top_raw_logits.tolist(), top_raw_indices.tolist())):
                         token_text = decode_token(idx)
-                        logits_entries.append({
+                        cos, row_norm = _cos_for(idx)
+                        entry = {
                             "token": token_text,
                             "token_id": idx,
                             "logit": logit_val,
-                            "rank": rank + 1
-                        })
+                            "rank": rank + 1,
+                        }
+                        if cos is not None:
+                            entry["cosine_sim"] = round(cos, 4)
+                        if row_norm is not None:
+                            entry["unemb_row_norm"] = round(row_norm, 4)
+                        logits_entries.append(entry)
 
                     # Greedy token (argmax of raw logits, before any sampling)
                     greedy_token_id = torch.argmax(raw_logits).item()
@@ -3177,13 +3519,21 @@ async def analyze_research_attention_stream(request: Dict[str, Any], authenticat
                         })
                         # Also add to logits if not present
                         if next_token_id not in [e["token_id"] for e in logits_entries]:
-                            logits_entries.append({
+                            cos_outlier, row_norm_outlier = _cos_for(next_token_id)
+                            outlier_entry = {
                                 "token": next_token_text,
                                 "token_id": next_token_id,
                                 "logit": selected_logit,
                                 "rank": selected_rank,
-                                "is_selected_outlier": True
-                            })
+                                "is_selected_outlier": True,
+                            }
+                            if cos_outlier is not None:
+                                outlier_entry["cosine_sim"] = round(cos_outlier, 4)
+                            if row_norm_outlier is not None:
+                                outlier_entry["unemb_row_norm"] = round(
+                                    row_norm_outlier, 4
+                                )
+                            logits_entries.append(outlier_entry)
 
                     # Build sampling metadata
                     eligible_count = int((probs_filtered > 0).sum().item())
@@ -3585,9 +3935,30 @@ async def analyze_research_attention_stream(request: Dict[str, Any], authenticat
                                 layer_entry["ffn_contribution"] = round(mlp_n / total, 4)
                         if layer_idx in gate_activation_stats:
                             layer_entry["gate_stats"] = gate_activation_stats[layer_idx]
-
-                        # Phase 5: Logit lens at sampled layers (every 8th layer)
-                        logit_lens_stride = max(1, n_layers // 5)
+                        if layer_idx in ffn_top_neurons:
+                            # `ffn_top_neurons` holds top-K SwiGLU
+                            # intermediate activations at the
+                            # predicting position — see
+                            # make_ffn_top_neurons_hook for the
+                            # per-neuron value semantics. The frontend
+                            # uses this to render Step 5's per-neuron
+                            # bar chart inside "Behind the SwiGLU"
+                            # without the backend needing a separate
+                            # endpoint or a 14,336-d cache entry.
+                            layer_entry["ffn_top_neurons"] = ffn_top_neurons[layer_idx]
+
+                        # Phase 5: Logit lens at every layer.
+                        # The lens projection is a single matmul
+                        # [1, hidden] × [hidden, vocab] ≈ 0.5-1ms on
+                        # Apple Silicon MPS — 40 layers × N_steps adds
+                        # only ~5-10s to a typical analyse run. Worth
+                        # the cost: every-layer sampling eliminates
+                        # the (L17, L24]-style uncertainty windows the
+                        # frontend's crystallisation narrative would
+                        # otherwise have to qualify, giving developers
+                        # a layer-precise commit point for trust
+                        # decisions about the model's prediction.
+                        logit_lens_stride = 1
                         if layer_idx % logit_lens_stride == 0 or layer_idx == n_layers - 1:
                             try:
                                 hidden_for_lens = current_hidden[-1].unsqueeze(0)  # [1, hidden_dim]
@@ -3866,6 +4237,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
@@ -3887,6 +4274,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,
@@ -3934,46 +4322,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")