Update modeling_neollm.py

modeling_neollm.py

@@ -20,7 +20,17 @@ Attention stack (orthogonal, all active simultaneously when enabled):
        moving-average bias that prevents collapse. Reduces first-token bias,
        increases attention entropy, and is complementary to Gated Attention
        (Bae et al. 2026, Table 2: Affine-Scaled + Gated > either alone).
-       Only active in eager attention mode (flash kernels lack weight access).
+
+       Flash/SDPA path (exact, no approximation):
+           Expanding [α·softmax(QKᵀ)+β]V distributively yields two terms:
+               α · flash_attn(Q,K,V)    — Flash computes this directly
+             + β · Σ_{j≤i} V_j          — causal prefix-sum of V (V.cumsum)
+           The V_cumsum tensor [B,H,S,d_head] is the only extra memory vs a
+           standard flash call. Semantically identical to the eager path;
+           per-weight softmax tensors (attn_weights_pre/post_affine) are
+           unavailable in flash mode and remain None in AnalysisState.
+
+       Eager path: exact with full weight access, used for interpretability.
 
 References:
     FANformer: "FANformer: Improving Large Language Models Through Effective
@@ -42,7 +52,8 @@ References:
 """
 
 import math
-from typing import Callable, Optional, Union, Tuple
+from dataclasses import dataclass
+from typing import Callable, List, Optional, Union, Tuple
 
 import torch
 import torch.nn.functional as F
@@ -60,14 +71,277 @@ from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_u
 from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel
 from transformers.processing_utils import Unpack
 from transformers.utils import TransformersKwargs, logging
-from .configuration_neollm import NeoLLMConfig
+from configuration_neollm import NeoLLMConfig
 
 from transformers import AutoConfig, AutoModel, AutoModelForCausalLM
 
 logger = logging.get_logger(__name__)
 
 
-# ==================== LEARNABLE MULTIPLIERS ====================
+# ==================== NATIVE ANALYSIS STATE ====================
+# All dataclasses below define the analysis container hierarchy.
+# This infrastructure is ONLY active when:
+#   1. model.enable_analysis() has been called, AND
+#   2. the model is in eval mode (model.eval() / not model.training)
+#
+# During training, analysis_state is always None — zero overhead,
+# zero interference with gradients, zero change to the training flow.
+#
+# Access after any inference call:
+#   state = model.last_analysis   # AnalysisState | None
+#
+# Every tensor stored here is detached from the computation graph.
+# Fields are None when the corresponding config flag is inactive,
+# so adding or removing a component only requires updating its own
+# dataclass — all other analysis paths remain untouched.
+
+
+@dataclass
+class FANAnalysis:
+    """
+    Decomposed output of a FANLayer call.
+    FANLayer'(X) = [cos(Wp·X) ‖ sin(Wp·X) ‖ (Wp̄·X + Bp̄)]
+    """
+    cosine_component: Optional[torch.Tensor] = None   # [*, p_output_dim]
+    sine_component:   Optional[torch.Tensor] = None   # [*, p_output_dim]
+    linear_component: Optional[torch.Tensor] = None   # [*, g_output_dim]
+
+
+@dataclass
+class SeeDNormAnalysis:
+    """
+    Internals of a SeeDNorm forward pass.
+    SeeDNorm(x) = [σ(x·β^T)·α + γ] ⊙ x/RMS(x)
+    """
+    rescale_factor: Optional[torch.Tensor] = None  # tanh(Σ x·β) — the dynamic gate
+    dynamic_scale:  Optional[torch.Tensor] = None  # rescale_factor·α + γ
+    x_normalized:   Optional[torch.Tensor] = None  # x / RMS(x) before dynamic_scale
+    output:         Optional[torch.Tensor] = None  # final output
+
+
+@dataclass
+class GPASAnalysis:
+    """
+    Internals of a GPAS forward pass.
+    GPAS(x) = x - silu(α)·x_detached
+    """
+    silu_alpha:           Optional[torch.Tensor] = None  # F.silu(self.alpha) — scalar
+    subtracted_component: Optional[torch.Tensor] = None  # silu(α) · x.detach()
+
+
+@dataclass
+class PolyNormAnalysis:
+    """
+    Internals of a PolyNorm forward pass (used as MLP act_fn).
+    Branches: x1 (linear), x2 (quadratic), x3 (cubic), each normalized.
+    x2 and x3 are partially orthogonalized against x1 via learned α.
+    """
+    x1:                Optional[torch.Tensor] = None  # RMS-normalized linear branch
+    x2_pre_exclusive:  Optional[torch.Tensor] = None  # quadratic branch before exclusivity
+    x3_pre_exclusive:  Optional[torch.Tensor] = None  # cubic branch before exclusivity
+    x2_post_exclusive: Optional[torch.Tensor] = None  # quadratic after exclusivity + renorm
+    x3_post_exclusive: Optional[torch.Tensor] = None  # cubic after exclusivity + renorm
+    alpha2:            Optional[torch.Tensor] = None  # sigmoid(exclusive_logits[0]) for x2
+    alpha3:            Optional[torch.Tensor] = None  # sigmoid(exclusive_logits[1]) for x3
+    weights:           Optional[torch.Tensor] = None  # self.weight [3] — branch mixing
+    bias:              Optional[torch.Tensor] = None  # self.bias scalar
+    output:            Optional[torch.Tensor] = None  # final PolyNorm output
+
+
+@dataclass
+class AttentionAnalysis:
+    """
+    Full attention internals for one NeoLLMAttention forward pass.
+    Fields present unconditionally capture the always-active path.
+    Fields guarded by config flags are None when that flag is off.
+    """
+    # ── FANLayer (always active) ──────────────────────────────────────
+    fan: Optional[FANAnalysis] = None                  # FAN components for attention
+
+    # ── Q/K/V projection (always active) ─────────────────────────────
+    q_raw:        Optional[torch.Tensor] = None  # Q before q_norm [B,S,H,d]
+    gate_raw:     Optional[torch.Tensor] = None  # gate chunk before sigmoid [B,S,H*d]
+    gate_sigmoid: Optional[torch.Tensor] = None  # sigmoid(gate) — Gated Attention weight
+    q_post_norm:  Optional[torch.Tensor] = None  # Q after SeeDNorm q_norm [B,H,S,d]
+    k_post_norm:  Optional[torch.Tensor] = None  # K after SeeDNorm k_norm [B,H,S,d]
+    v_raw:        Optional[torch.Tensor] = None  # V raw (pre MEA/LUCID) [B,H,S,d]
+
+    # ── RoPE (always active) ──────────────────────────────────────────
+    q_post_rope: Optional[torch.Tensor] = None   # Q after RoPE [B,H,S,d]
+    k_post_rope: Optional[torch.Tensor] = None   # K after RoPE [B,H,S,d]
+
+    # ── Momentum (conditional on use_momentum_attention) ──────────────
+    q_momentum_delta: Optional[torch.Tensor] = None  # causal_first_difference(Q)
+    k_momentum_delta: Optional[torch.Tensor] = None  # causal_first_difference(K)
+    q_post_momentum:  Optional[torch.Tensor] = None  # Q + γ·Δ
+    k_post_momentum:  Optional[torch.Tensor] = None  # K + γ·Δ
+
+    # ── MEA head mixing (conditional on use_mea_attention) ────────────
+    mea_key_mix_matrix:   Optional[torch.Tensor] = None  # mea_key_mix [H_comp,H_kv]
+    mea_value_mix_matrix: Optional[torch.Tensor] = None  # mea_value_mix [H_comp,H_kv]
+    k_post_mea:           Optional[torch.Tensor] = None  # K after head mixing
+    v_post_mea:           Optional[torch.Tensor] = None  # V after head mixing
+
+    # ── LUCID preconditioner (conditional on use_lucid_attention) ─────
+    lucid_preconditioner: Optional[torch.Tensor] = None  # lower-triangular prec matrix
+    v_post_lucid:         Optional[torch.Tensor] = None  # V after triangular solve
+
+    # ── Affine-Scaled Attention (conditional on use_affine_scaled_attention) ──
+    alpha_per_head:           Optional[torch.Tensor] = None  # α [B,H,S,1] in [0,1]
+    beta_per_head:            Optional[torch.Tensor] = None  # β [B,H,S,1] moving-avg bias
+    alpha_moving_avg:         Optional[torch.Tensor] = None  # alpha_ma EMA snapshot
+    attn_weights_pre_affine:  Optional[torch.Tensor] = None  # softmax weights before α,β
+    attn_weights_post_affine: Optional[torch.Tensor] = None  # α·softmax + β
+
+    # ── Standard attention weights (eager non-affine path) ────────────
+    attn_weights: Optional[torch.Tensor] = None  # softmax weights (None for flash/sdpa)
+
+    # ── Post-SDPA (always active) ─────────────────────────────────────
+    attn_output_raw: Optional[torch.Tensor] = None  # SDPA output [B,S,H,d]
+
+    # ── MEA output norm (conditional on use_mea_attention) ────────────
+    attn_output_post_mea_norm: Optional[torch.Tensor] = None  # after MEAHeadSeeDNorm
+
+    # ── XSA (conditional on use_xsa) ─────────────────────────────────
+    xsa_self_position_component: Optional[torch.Tensor] = None  # proj subtracted
+    attn_output_post_xsa:        Optional[torch.Tensor] = None  # after XSA removal
+
+    # ── Directional Routing (conditional on use_directional_routing) ──
+    direction_vecs_normalized: Optional[torch.Tensor] = None  # unit-norm d [H,K,d]
+    dr_router_logits:          Optional[torch.Tensor] = None  # MLP router output [B,H*K]
+    dr_routing_weights:        Optional[torch.Tensor] = None  # sigmoid(T·logits) [B,H,K]
+    dr_projection:             Optional[torch.Tensor] = None  # (o·d) scalars [B,S,H,K]
+    dr_suppression:            Optional[torch.Tensor] = None  # Σ r·proj·d [B,S,H,d]
+    attn_output_post_routing:  Optional[torch.Tensor] = None  # after DR removal
+
+    # ── Gate and o_proj (always active) ───────────────────────────────
+    attn_output_pre_gate: Optional[torch.Tensor] = None  # pre gate multiply [B,S,H,d]
+    attn_output_final:    Optional[torch.Tensor] = None  # after o_proj [B,S,D]
+
+
+@dataclass
+class MLPAnalysis:
+    """
+    Internals of a NeoLLMMLP forward pass.
+    SwiGLU-like: down_proj(dropout(PolyNorm(gate_proj(fan)) · up_proj(fan)))
+    """
+    fan:              Optional[FANAnalysis] = None       # FAN components for MLP
+    gate_proj_output: Optional[torch.Tensor] = None     # gate_proj(x_fan) [B,S,I]
+    up_proj_output:   Optional[torch.Tensor] = None     # up_proj(x_fan) [B,S,I]
+    polynorm:         Optional[PolyNormAnalysis] = None  # PolyNorm of gate branch
+    act_times_up:     Optional[torch.Tensor] = None     # PolyNorm(gate)·up [B,S,I]
+    output:           Optional[torch.Tensor] = None     # after down_proj [B,S,D]
+
+
+@dataclass
+class JTokMAnalysis:
+    """
+    Internals of a LeviathanJTokM forward pass for one decoder layer.
+    """
+    surfaces:         Optional[torch.Tensor] = None  # all n_e surface outputs [N,n_e,D]
+    router_logits:    Optional[torch.Tensor] = None  # pre-TopK logits [N,n_e]
+    topk_indices:     Optional[torch.Tensor] = None  # selected surface indices [N,K]
+    routing_weights:  Optional[torch.Tensor] = None  # normalized sigmoid weights [N,K]
+    mixed_pre_norm:   Optional[torch.Tensor] = None  # weighted sum before norm [N,D]
+    mixed_normalized: Optional[torch.Tensor] = None  # direction-normalized mixed [N,D]
+    delta_r:          Optional[torch.Tensor] = None  # final injection (scaled) [N,D]
+    p_sum:            Optional[torch.Tensor] = None  # routing probability sum [n_e]
+    f_sum:            Optional[torch.Tensor] = None  # load fraction sum [n_e]
+    lns_scale:        Optional[float] = None          # 1/√(2ℓ) scaling factor
+
+
+@dataclass
+class AttnResAnalysis:
+    """
+    Depth-wise softmax attention weights from AttnRes sublayer calls.
+    Only populated when use_attn_res=True.
+    """
+    weights_pre_attn: Optional[torch.Tensor] = None  # softmax over sources [N+1,B,S]
+    weights_pre_mlp:  Optional[torch.Tensor] = None  # softmax over sources [N+1,B,S]
+    sources_count:    Optional[int] = None            # number of sources including partial
+
+
+@dataclass
+class LayerAnalysis:
+    """
+    Complete analysis snapshot for one NeoLLMDecoderLayer forward pass.
+    Sub-objects are None when the corresponding config flag is inactive.
+    """
+    layer_idx: int = 0
+
+    # Hidden state boundaries
+    hidden_states_input:  Optional[torch.Tensor] = None  # entering this layer
+    hidden_states_output: Optional[torch.Tensor] = None  # leaving this layer
+    h_tilde:              Optional[torch.Tensor] = None  # post-attn residual, pre-MLP
+
+    # Attention sublayer
+    seednorm_pre_attn:  Optional[SeeDNormAnalysis] = None  # input_layernorm
+    lns_attn_output:    Optional[torch.Tensor] = None      # after LNS(1/√ℓ)
+    attention:          Optional[AttentionAnalysis] = None  # full attention analysis
+    attn_contribution:  Optional[torch.Tensor] = None      # attn output before residual
+    gpas_attn:          Optional[GPASAnalysis] = None       # GPAS after attn residual
+
+    # MLP sublayer
+    seednorm_post_attn: Optional[SeeDNormAnalysis] = None  # post_attention_layernorm
+    lns_mlp_output:     Optional[torch.Tensor] = None      # after LNS(1/√ℓ)
+    mlp:                Optional[MLPAnalysis] = None       # full MLP analysis
+    mlp_contribution:   Optional[torch.Tensor] = None      # MLP output before residual
+    gpas_mlp:           Optional[GPASAnalysis] = None      # GPAS after MLP residual
+
+    # Optional components (None when inactive)
+    jtokm:    Optional[JTokMAnalysis] = None    # if use_jtokm
+    attn_res: Optional[AttnResAnalysis] = None  # if use_attn_res
+
+
+@dataclass
+class GeneratorAnalysis:
+    """
+    Internals of a LeviathanGenerator forward pass.
+    Only populated when use_token_generator=True.
+    """
+    z_raw:              Optional[torch.Tensor] = None  # [N,d_seed] codebook sum
+    z_tilde:            Optional[torch.Tensor] = None  # [N,d_seed] JTok-M path output
+    B_vals:             Optional[torch.Tensor] = None  # [N,d_seed,n_knots] B-spline basis
+    z_all_pre_norm:     Optional[torch.Tensor] = None  # [N,M,d_seed] per-head pre-sigmoid
+    z_all_post_sigmoid: Optional[torch.Tensor] = None  # [N,M,d_seed] per-head post-sigmoid
+    modes_all:          Optional[torch.Tensor] = None  # [N,M,krank] KHRONOS tensor product
+    embeddings:         Optional[torch.Tensor] = None  # [N,hidden_size] final output
+
+
+@dataclass
+class AnalysisState:
+    """
+    Root analysis container for one NeoLLMForCausalLM forward pass.
+
+    Populated when model.enable_analysis() is active AND model.training=False.
+    All tensors are detached from the computation graph.
+
+    Structure:
+        .input_ids              — original input token ids
+        .embeddings             — token embeddings entering the decoder stack
+        .final_hidden_states    — post-norm output entering lm_head
+        .generator              — LeviathanGenerator internals (use_token_generator only)
+        .layers[i]              — per-layer LayerAnalysis, indexed by layer_idx
+        .jtokm_aux_stats        — raw per-layer load-balancing tuples (use_jtokm only)
+        .attn_res_sources_final — AttnRes source list at end of forward (use_attn_res only)
+        .logits                 — lm_head output (only when labels is None)
+
+    Access pattern:
+        model.eval()
+        model.enable_analysis()
+        _ = model(input_ids)
+        state = model.last_analysis
+        alpha = state.layers[3].attention.alpha_per_head
+    """
+    input_ids:              Optional[torch.Tensor] = None
+    embeddings:             Optional[torch.Tensor] = None
+    final_hidden_states:    Optional[torch.Tensor] = None
+
+    generator:              Optional[GeneratorAnalysis] = None
+    layers:                 Optional[List[LayerAnalysis]] = None
+    jtokm_aux_stats:        Optional[list] = None
+    attn_res_sources_final: Optional[list] = None
+    logits:                 Optional[torch.Tensor] = None
 
 class ScalarMultiplier(nn.Module):
     """
@@ -454,6 +728,7 @@ class LeviathanGenerator(nn.Module):
         self,
         token_ids: torch.Tensor,
         return_internals: bool = False,
+        analysis: Optional[GeneratorAnalysis] = None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]:
         """
         Generate embeddings from discrete token indices.
@@ -478,6 +753,9 @@ class LeviathanGenerator(nn.Module):
             token_ids:        (batch, seq_len) or (seq_len,)
             return_internals: if True, also return z_tilde and B_vals for
                               reuse by JTok-M surfaces in every decoder layer.
+            analysis:         Optional[GeneratorAnalysis] — when not None and model
+                              is in eval mode with analysis armed, deposits all
+                              internal tensors (detached). No-op during training.
         Returns:
             embeddings [*token_ids.shape, hidden_size],
             or (embeddings, z_tilde [N, d_seed], B_vals [N, d_seed, n_knots])
@@ -494,12 +772,19 @@ class LeviathanGenerator(nn.Module):
         for r in range(self.k):
             z = z + self.codebooks[r][coords_flat[:, r]]
 
+        if analysis is not None:
+            analysis.z_raw = z.detach()
+
         # ── JTok-M shared path ────────────────────────────────────────────
         # Produces z_tilde and B_vals consumed by every decoder layer's
         # JTok-M module. This path is unchanged and uses fixed scalar scale.
         z_tilde = torch.sigmoid(self.seed_norm(self.seed_proj(z)))  # [N, d_seed]
         B_vals  = self._bspline_basis(z_tilde.clamp(0.0, 1.0))     # [N, d_seed, n_knots]
 
+        if analysis is not None:
+            analysis.z_tilde = z_tilde.detach()
+            analysis.B_vals  = B_vals.detach()
+
         # ── Per-head generator path (fully vectorized, 6 kernels) ────────
         # All 8 heads are processed simultaneously. No Python loop.
         # Maximum intermediate tensor [N, M, d_seed, n_knots] appears once.
@@ -509,6 +794,9 @@ class LeviathanGenerator(nn.Module):
         z_all = F.linear(z.to(target_dtype), self.head_proj_weight)
         z_all = z_all.view(N, self.num_modes, self.d_seed)   # [N, M, d_seed]
 
+        if analysis is not None:
+            analysis.z_all_pre_norm = z_all.detach()
+
         # Kernel 2: per-head LayerNorm + sigmoid(x/2)
         # Manual LN over last dim with independent weight/bias per head.
         # Mathematically identical to 8 separate nn.LayerNorm(d_seed).
@@ -520,6 +808,9 @@ class LeviathanGenerator(nn.Module):
                       + self.head_norm_bias.unsqueeze(0)
         z_all = torch.sigmoid(z_all / 2.0)                   # [N, M, d_seed]
 
+        if analysis is not None:
+            analysis.z_all_post_sigmoid = z_all.detach()
+
         # Kernel 3: vectorized B-spline basis for all heads
         # head_scale [M, d_seed] is used inside _bspline_basis_all_heads
         B_all = self._bspline_basis_all_heads(
@@ -529,6 +820,9 @@ class LeviathanGenerator(nn.Module):
         # Kernel 4: vectorized KHRONOS tensor product for all heads
         modes_all = self._khronos_all_heads(B_all)            # [N, M, krank]
 
+        if analysis is not None:
+            analysis.modes_all = modes_all.detach()
+
         # Kernel 5: project all heads to hidden_size and sum
         # einsum: token n, head m, krank k → hidden d (summed over m)
         # head_out_weight [M, krank, hidden_size]
@@ -541,6 +835,9 @@ class LeviathanGenerator(nn.Module):
         # No W_res — confirmed absent in the authors' implementation
         e = e.reshape(*orig_shape, self.hidden_size)
 
+        if analysis is not None:
+            analysis.embeddings = e.detach()
+
         if return_internals:
             return e, z_tilde, B_vals
         return e
@@ -627,6 +924,7 @@ class LeviathanJTokM(nn.Module):
         B_vals: torch.Tensor,
         z_tilde: torch.Tensor,
         target_dtype: torch.dtype,
+        analysis: Optional[JTokMAnalysis] = None,
     ) -> torch.Tensor:
         """
         Evaluate all n_e surfaces vectorized over the full token batch.
@@ -638,9 +936,10 @@ class LeviathanJTokM(nn.Module):
         All shapes are static → torch.compile compatible.
 
         Args:
-            B_vals:  [N, d_seed, n_knots]  float32
-            z_tilde: [N, d_seed]
+            B_vals:       [N, d_seed, n_knots]  float32
+            z_tilde:      [N, d_seed]
             target_dtype: model dtype
+            analysis:     Optional[JTokMAnalysis] — deposits surfaces when not None.
         Returns:
             surfaces: [N, n_e, D]
         """
@@ -669,7 +968,12 @@ class LeviathanJTokM(nn.Module):
         z = z_tilde.to(target_dtype)
         out_res   = torch.einsum("nd,idc->nic", z, self.W_res.to(target_dtype))
 
-        return out_modes + out_res  # [N, n_e, D]
+        surfaces = out_modes + out_res  # [N, n_e, D]
+
+        if analysis is not None:
+            analysis.surfaces = surfaces.detach()
+
+        return surfaces
 
     # ── Router ────────────────────────────────────────────────────────────
 
@@ -680,6 +984,7 @@ class LeviathanJTokM(nn.Module):
         self,
         h_tilde: torch.Tensor,
         surfaces: torch.Tensor,
+        analysis: Optional[JTokMAnalysis] = None,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         """
         Context-dependent routing over h_tilde (hidden state after attention).
@@ -698,6 +1003,7 @@ class LeviathanJTokM(nn.Module):
         Args:
             h_tilde:  [N, D] — hidden state after attention (before MLP)
             surfaces: [N, n_e, D]
+            analysis: Optional[JTokMAnalysis] — deposits routing data when not None.
         Returns:
             mixed: [N, D]
             aux_stats: (p_sum [n_e], f_sum [n_e], N) for loss accumulation
@@ -720,6 +1026,12 @@ class LeviathanJTokM(nn.Module):
         selected = surfaces.gather(dim=1, index=idx_exp)              # [N, K, D]
         mixed    = (w.unsqueeze(-1) * selected).sum(dim=1)           # [N, D]
 
+        if analysis is not None:
+            analysis.router_logits   = g.detach()
+            analysis.topk_indices    = topk_idx.detach()
+            analysis.routing_weights = w.detach()
+            analysis.mixed_pre_norm  = mixed.detach()
+
         # Load-balancing statistics for aux loss (Appendix B, Yang et al. 2026)
         # p_i = mean routing probability over batch
         # f_i = fraction of tokens actually routed to i
@@ -741,31 +1053,41 @@ class LeviathanJTokM(nn.Module):
         h_tilde: torch.Tensor,
         z_tilde: torch.Tensor,
         B_vals: torch.Tensor,
+        analysis: Optional[JTokMAnalysis] = None,
     ) -> Tuple[torch.Tensor, Tuple]:
         """
         Compute additive JTok-M residual for one decoder layer.
 
         Args:
-            h_tilde: [N, D]             hidden state after attention (before MLP)
-            z_tilde: [N, d_seed]        latent coordinate from generator
-            B_vals:  [N, d_seed, n_k]   B-spline basis (computed once, reused)
+            h_tilde:  [N, D]           hidden state after attention (before MLP)
+            z_tilde:  [N, d_seed]      latent coordinate from generator
+            B_vals:   [N, d_seed, n_k] B-spline basis (computed once, reused)
+            analysis: Optional[JTokMAnalysis] — deposits all JTok-M internals when
+                      not None. No-op during training (analysis is always None then).
         Returns:
-            delta_r: [N, D]             additive residual (already scaled)
+            delta_r:   [N, D]          additive residual (already scaled)
             aux_stats: tuple for accumulating load-balance loss
         """
         target_dtype = h_tilde.dtype
 
         # All n_e surfaces in one vectorized pass
-        surfaces = self._eval_surfaces(B_vals, z_tilde, target_dtype)  # [N, n_e, D]
+        surfaces = self._eval_surfaces(B_vals, z_tilde, target_dtype, analysis=analysis)
 
         # Context-dependent routing
-        mixed, aux_stats = self._route_and_mix(h_tilde, surfaces)      # [N, D]
+        mixed, aux_stats = self._route_and_mix(h_tilde, surfaces, analysis=analysis)
 
         # Normalise direction, apply scaler, scale with 1/√(2ℓ)
         # Norm_ε decouples direction from magnitude (JTok Appendix D.2)
         mixed_norm = mixed / (mixed.norm(dim=-1, keepdim=True) + self.norm_eps)
         delta_r    = self.lns_scale * self.scaler * mixed_norm          # [N, D]
 
+        if analysis is not None:
+            analysis.mixed_normalized = mixed_norm.detach()
+            analysis.delta_r          = delta_r.detach()
+            analysis.p_sum            = aux_stats[0].detach()
+            analysis.f_sum            = aux_stats[1].detach()
+            analysis.lns_scale        = self.lns_scale
+
         return delta_r, aux_stats
 
 
@@ -825,10 +1147,20 @@ class FANLayer(nn.Module):
         if self.input_linear.bias is not None:
             nn.init.zeros_(self.input_linear.bias)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self,
+        x: torch.Tensor,
+        analysis: Optional[FANAnalysis] = None,
+    ) -> torch.Tensor:
         pg = self.input_linear(x)
         p, g = torch.split(pg, [self.p_output_dim, self.g_output_dim], dim=-1)
-        return torch.cat([torch.cos(p), torch.sin(p), g], dim=-1)
+        cos_p = torch.cos(p)
+        sin_p = torch.sin(p)
+        if analysis is not None:
+            analysis.cosine_component = cos_p.detach()
+            analysis.sine_component   = sin_p.detach()
+            analysis.linear_component = g.detach()
+        return torch.cat([cos_p, sin_p, g], dim=-1)
 
 
 class LNS(nn.Module):
@@ -853,8 +1185,17 @@ class GPAS(nn.Module):
         self.d_model = d_model
         self.alpha   = nn.Parameter(torch.zeros(1))
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return x - F.silu(self.alpha) * x.detach()
+    def forward(
+        self,
+        x: torch.Tensor,
+        analysis: Optional[GPASAnalysis] = None,
+    ) -> torch.Tensor:
+        silu_alpha = F.silu(self.alpha)
+        subtracted = silu_alpha * x.detach()
+        if analysis is not None:
+            analysis.silu_alpha           = silu_alpha.detach()
+            analysis.subtracted_component = subtracted.detach()
+        return x - subtracted
 
 
 class SeeDNorm(nn.Module):
@@ -883,7 +1224,11 @@ class SeeDNorm(nn.Module):
     def _rms_norm(self, x: torch.Tensor) -> torch.Tensor:
         return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
 
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self,
+        x: torch.Tensor,
+        analysis: Optional[SeeDNormAnalysis] = None,
+    ) -> torch.Tensor:
         x_for_dynamic  = F.dropout(x, p=self.dropout_input)
         rescale_factor = torch.tanh(
             torch.sum(x_for_dynamic * self.beta, dim=-1, keepdim=True)
@@ -891,7 +1236,13 @@ class SeeDNorm(nn.Module):
         dynamic_scale  = rescale_factor * self.alpha + self.gamma
         x_normalized   = self._rms_norm(x.float())
         x_normalized   = F.dropout(x_normalized, p=self.dropout_hidden)
-        return (x_normalized * dynamic_scale.float()).type_as(x)
+        output = (x_normalized * dynamic_scale.float()).type_as(x)
+        if analysis is not None:
+            analysis.rescale_factor = rescale_factor.detach()
+            analysis.dynamic_scale  = dynamic_scale.detach()
+            analysis.x_normalized   = x_normalized.detach()
+            analysis.output         = output.detach()
+        return output
 
     def extra_repr(self) -> str:
         return (f"dim={self.dim}, eps={self.eps}, "
@@ -1107,6 +1458,7 @@ def affine_scaled_eager_attention_forward(
     alpha: torch.Tensor,
     beta: torch.Tensor,
     dropout: float = 0.0,
+    attn_analysis: Optional[AttentionAnalysis] = None,
     **kwargs: Unpack[TransformersKwargs],
 ):
     """
@@ -1129,8 +1481,10 @@ def affine_scaled_eager_attention_forward(
     Reference: Bae et al. (2026), Affine-Scaled Attention, Eq. 6–8.
 
     Args:
-        alpha: [batch, num_heads, seq_q, 1]  — input-dependent scale per query
-        beta:  [batch, num_heads, seq_q, 1]  — input-dependent bias per query
+        alpha:        [batch, num_heads, seq_q, 1]  — input-dependent scale per query
+        beta:         [batch, num_heads, seq_q, 1]  — input-dependent bias per query
+        attn_analysis: Optional[AttentionAnalysis]  — deposits pre/post-affine weights
+                       when not None. No-op during training (always None then).
     """
     key_states   = repeat_kv(key,   module.num_key_value_groups)
     value_states = repeat_kv(value, module.num_key_value_groups)
@@ -1139,19 +1493,137 @@ def affine_scaled_eager_attention_forward(
     if attention_mask is not None:
         attn_weights = attn_weights + attention_mask[:, :, :, : key_states.shape[-2]]
 
-    attn_weights = nn.functional.softmax(
+    attn_weights_softmax = nn.functional.softmax(
         attn_weights, dim=-1, dtype=torch.float32
     ).to(query.dtype)
 
+    if attn_analysis is not None:
+        attn_analysis.attn_weights_pre_affine = attn_weights_softmax.detach()
+
     # Affine reweighting: relaxes the unit-sum constraint of softmax.
     # α ∈ [0, 1] scales down the softmax distribution (input-adaptively per head).
     # β offsets to prevent collapse when α deviates from its running mean.
     # Shapes: α, β are [B, H, S_q, 1], attn_weights is [B, H, S_q, S_k] → broadcast.
-    attn_weights = alpha * attn_weights + beta
+    attn_weights_affine = alpha * attn_weights_softmax + beta
 
-    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)
-    attn_output  = torch.matmul(attn_weights, value_states).transpose(1, 2).contiguous()
-    return attn_output, attn_weights
+    if attn_analysis is not None:
+        attn_analysis.attn_weights_post_affine = attn_weights_affine.detach()
+
+    attn_weights_affine = nn.functional.dropout(
+        attn_weights_affine, p=dropout, training=module.training
+    )
+    attn_output  = torch.matmul(attn_weights_affine, value_states).transpose(1, 2).contiguous()
+    return attn_output, attn_weights_affine
+
+
+def affine_scaled_flash_attention_forward(
+    module: nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: Optional[torch.Tensor],
+    scaling: float,
+    alpha: torch.Tensor,
+    beta: torch.Tensor,
+    dropout: float = 0.0,
+    attn_analysis: Optional[AttentionAnalysis] = None,
+    **kwargs: Unpack[TransformersKwargs],
+):
+    """
+    Affine-Scaled Attention — flash/sdpa path.
+
+    Exact mathematical decomposition of [α·softmax(QKᵀ)+β]V using only the
+    public flash/sdpa interface — no kernel modification required.
+
+    Derivation
+    ----------
+    The paper formula expands distributively:
+
+        [α · softmax(QKᵀ/√dk) + β] · V
+        = α · [softmax(QKᵀ/√dk) · V]   ← term 1: standard flash output
+        + β · [Σ_{j≤i} V_j]             ← term 2: causal prefix-sum of V
+
+    Term 2 follows because β is a scalar per query (broadcast over all S_k),
+    and softmax weights sum to 1 over the causal window:
+        β_i · Σ_j w_{i,j} · V_j = β_i · Σ_{j≤i} V_j  (since Σ w_{i,j} = 1)
+
+    Dropout
+    -------
+    The eager path drops entries of the combined weight matrix (α·softmax + β)
+    before multiplying by V. With the flash interface we cannot access that
+    combined matrix, so we apply dropout=0 to the flash kernel and instead
+    apply nn.functional.dropout to the final combined output tensor. This is
+    output dropout rather than weight dropout — a different (but standard)
+    regularisation that achieves the same intent without the intermediate
+    weight matrix. During inference dropout=0 so the paths are identical.
+
+    Padding mask
+    ------------
+    The V_cumsum must not accumulate values from padding positions, since the
+    flash kernel zeros those out internally but a plain cumsum does not.
+    We zero V at padding positions before the cumsum by reading the diagonal
+    of the attention_mask (valid positions have mask value 0, padding has -inf).
+
+    Memory overhead vs standard flash call
+    ---------------------------------------
+    One extra tensor of shape [B, H_q, S, d_head] for V_cumsum.
+    At (B=1, H=8, S=2048, d_head=64, bf16): ≈ 4 MB per call, ≈ 48 MB total
+    across 12 layers. Allocated and freed within each forward call.
+
+    Attention-weight analysis fields (attn_weights_pre_affine,
+    attn_weights_post_affine) remain None — those tensors are never
+    materialised by the flash kernel and cannot be recovered.
+    All other AnalysisState fields (alpha, beta, alpha_ma) are deposited
+    by the caller before this function is invoked.
+
+    Args:
+        alpha: [B, H_q, S_q, 1] — input-dependent scale per query, in [0, 1]
+        beta:  [B, H_q, S_q, 1] — moving-average bias per query
+        attn_analysis: deposited by caller; this function does not write to it.
+    """
+    # ── Term 1: standard flash / sdpa output ─────────────────────────────
+    # dropout=0.0: we apply dropout to the combined output below instead.
+    attn_fn   = ALL_ATTENTION_FUNCTIONS[module.config._attn_implementation]
+    flash_out, _ = attn_fn(
+        module, query, key, value, attention_mask,
+        dropout=0.0, scaling=scaling, **kwargs,
+    )
+    # flash_out: [B, S, H_q, d_head]  — HF wrappers all return this layout
+
+    # ── Term 2: β · causal prefix-sum of V ───────────────────────────────
+    # Expand V from KV heads to query heads for GQA.
+    value_expanded = repeat_kv(value, module.num_key_value_groups)  # [B, H_q, S, d_head]
+
+    # Zero out padded positions so they don't accumulate into the prefix sum.
+    # attention_mask is [B, 1, S, S] with 0 at valid positions, -inf at padding.
+    # The diagonal gives the per-position validity: valid → 0, padding → -inf.
+    if attention_mask is not None and attention_mask.ndim == 4:
+        diag = attention_mask.diagonal(dim1=-2, dim2=-1)         # [B, 1, S]
+        # valid=True where diag==0 (not -inf), padding=False
+        valid = (diag == 0).to(value_expanded.dtype)             # [B, 1, S]
+        valid = valid.unsqueeze(-1)                              # [B, 1, S, 1]
+        # broadcast over H_q and d_head; zero out V at padding positions
+        value_expanded = value_expanded * valid                  # [B, H_q, S, d_head]
+
+    # Causal prefix-sum: position i accumulates all valid j ≤ i.
+    v_cumsum   = value_expanded.cumsum(dim=2)                    # [B, H_q, S, d_head]
+    # Transpose to match flash_out layout.
+    v_cumsum_t = v_cumsum.transpose(1, 2).contiguous()          # [B, S, H_q, d_head]
+
+    # α, β: [B, H_q, S, 1] → [B, S, H_q, 1] to broadcast over d_head.
+    alpha_t = alpha.permute(0, 2, 1, 3)   # [B, S, H_q, 1]
+    beta_t  = beta.permute(0, 2, 1, 3)    # [B, S, H_q, 1]
+
+    # ── Combine and apply dropout to the full affine output ────���──────────
+    output = alpha_t * flash_out + beta_t * v_cumsum_t           # [B, S, H_q, d_head]
+
+    # Apply output dropout on the combined affine result.
+    # This regularises the full [α·flash + β·V_cumsum] output consistently.
+    if dropout > 0.0 and module.training:
+        output = nn.functional.dropout(output, p=dropout, training=True)
+
+    # attn_weights is None — flash never exposes the softmax weight matrix.
+    return output, None
 
 
 class NeoLLMAttention(nn.Module):
@@ -1336,19 +1808,49 @@ class NeoLLMAttention(nn.Module):
             self.direction_vecs   = None
             self.direction_router = None
 
-    def _apply_momentum_attention(self, q, k):
+    def _apply_momentum_attention(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        attn_analysis: Optional[AttentionAnalysis] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
         if not self.use_momentum_attention or self.momentum_gamma == 0.0:
             return q, k
-        return q + self.momentum_gamma * causal_first_difference(q), \
-               k + self.momentum_gamma * causal_first_difference(k)
-
-    def _apply_mea_head_mixing(self, k, v):
+        dq    = causal_first_difference(q)
+        dk    = causal_first_difference(k)
+        q_new = q + self.momentum_gamma * dq
+        k_new = k + self.momentum_gamma * dk
+        if attn_analysis is not None:
+            attn_analysis.q_momentum_delta = dq.detach()
+            attn_analysis.k_momentum_delta = dk.detach()
+            attn_analysis.q_post_momentum  = q_new.detach()
+            attn_analysis.k_post_momentum  = k_new.detach()
+        return q_new, k_new
+
+    def _apply_mea_head_mixing(
+        self,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        attn_analysis: Optional[AttentionAnalysis] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
         if not self.use_mea_attention:
             return k, v
-        return head_linear_compose(k, self.mea_key_mix).contiguous(), \
-               head_linear_compose(v, self.mea_value_mix).contiguous()
-
-    def _apply_lucid_preconditioner(self, k, v, attention_mask):
+        k_mixed = head_linear_compose(k, self.mea_key_mix).contiguous()
+        v_mixed = head_linear_compose(v, self.mea_value_mix).contiguous()
+        if attn_analysis is not None:
+            attn_analysis.mea_key_mix_matrix   = self.mea_key_mix.detach()
+            attn_analysis.mea_value_mix_matrix = self.mea_value_mix.detach()
+            attn_analysis.k_post_mea           = k_mixed.detach()
+            attn_analysis.v_post_mea           = v_mixed.detach()
+        return k_mixed, v_mixed
+
+    def _apply_lucid_preconditioner(
+        self,
+        k: torch.Tensor,
+        v: torch.Tensor,
+        attention_mask: Optional[torch.Tensor],
+        attn_analysis: Optional[AttentionAnalysis] = None,
+    ) -> torch.Tensor:
         if not self.use_lucid_attention:
             return v.contiguous()
         key_rn = rms_key_unit_norm(k, eps=self.lucid_attention_eps)
@@ -1360,15 +1862,20 @@ class NeoLLMAttention(nn.Module):
         eye  = torch.eye(prec.shape[-1], device=prec.device, dtype=prec.dtype).view(
             1, 1, prec.shape[-1], prec.shape[-1]
         )
-        prec = prec + eye * (1.0 - prec.diagonal(dim1=-2, dim2=-1).unsqueeze(-1))
-        return torch.linalg.solve_triangular(
+        prec   = prec + eye * (1.0 - prec.diagonal(dim1=-2, dim2=-1).unsqueeze(-1))
+        result = torch.linalg.solve_triangular(
             prec, v.float(), upper=False, unitriangular=True
         ).to(v.dtype).contiguous()
+        if attn_analysis is not None:
+            attn_analysis.lucid_preconditioner = prec.detach()
+            attn_analysis.v_post_lucid         = result.detach()
+        return result
 
     def _apply_directional_routing(
         self,
         attn_out: torch.Tensor,
         hidden_states: torch.Tensor,
+        attn_analysis: Optional[AttentionAnalysis] = None,
     ) -> torch.Tensor:
         """
         Directional suppression at position C (post-XSA, pre-reshape).
@@ -1377,6 +1884,7 @@ class NeoLLMAttention(nn.Module):
             attn_out:      [B, S, H, d_head] — output after XSA and SeeDNorm.
             hidden_states: [B, S, hidden_size] — pre-FAN residual stream,
                            used as router input (same as paper's x_i).
+            attn_analysis: Optional[AttentionAnalysis] — deposits DR internals.
         Returns:
             [B, S, H, d_head] with selected directional components suppressed.
         """
@@ -1416,8 +1924,17 @@ class NeoLLMAttention(nn.Module):
         # Σ_k weighted_{h,k} · d_{h,k}:
         # weighted [B, S, H, K] × d [H, K, D] → [B, S, H, D]
         suppression = torch.einsum("bshk,hkd->bshd", weighted, d)
+        result      = attn_out - suppression
+
+        if attn_analysis is not None:
+            attn_analysis.direction_vecs_normalized = d.detach()
+            attn_analysis.dr_router_logits          = logits.detach()
+            attn_analysis.dr_routing_weights        = r.squeeze(1).detach()
+            attn_analysis.dr_projection             = proj.detach()
+            attn_analysis.dr_suppression            = suppression.detach()
+            attn_analysis.attn_output_post_routing  = result.detach()
 
-        return attn_out - suppression
+        return result
 
     def forward(
         self,
@@ -1425,11 +1942,13 @@ class NeoLLMAttention(nn.Module):
         position_embeddings: tuple[torch.Tensor, torch.Tensor],
         attention_mask: Optional[torch.Tensor] = None,
         first_layer_fan: Optional[torch.Tensor] = None,
+        attn_analysis: Optional[AttentionAnalysis] = None,
         **kwargs: Unpack[FlashAttentionKwargs],
     ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
         input_shape = hidden_states.shape[:-1]
 
-        h_fan = self.fan_layer(hidden_states)
+        fan_a = attn_analysis.fan if attn_analysis is not None else None
+        h_fan = self.fan_layer(hidden_states, analysis=fan_a)
         if first_layer_fan is not None:
             h_fan = self.lambda_1 * first_layer_fan + self.lambda_2 * h_fan
         current_layer_fan = h_fan.clone()
@@ -1443,27 +1962,42 @@ class NeoLLMAttention(nn.Module):
         )
         gate = gate.reshape(*input_shape, -1)
 
+        if attn_analysis is not None:
+            attn_analysis.q_raw    = q_raw.detach()
+            attn_analysis.gate_raw = gate.detach()
+
         q = self.q_norm(q_raw.view(query_shape)).transpose(1, 2)
         k = self.k_norm(self.k_proj(h_fan).view(kv_shape)).transpose(1, 2)
         v = self.v_proj(h_fan).view(kv_shape).transpose(1, 2)
 
+        if attn_analysis is not None:
+            attn_analysis.q_post_norm = q.detach()
+            attn_analysis.k_post_norm = k.detach()
+            attn_analysis.v_raw       = v.detach()
+
         cos, sin = position_embeddings
         q, k     = apply_rotary_pos_emb(q, k, cos, sin)
-        q, k     = self._apply_momentum_attention(q, k)
-        k, v     = self._apply_mea_head_mixing(k, v)
-        v        = self._apply_lucid_preconditioner(k, v, attention_mask)
+
+        if attn_analysis is not None:
+            attn_analysis.q_post_rope = q.detach()
+            attn_analysis.k_post_rope = k.detach()
+
+        q, k = self._apply_momentum_attention(q, k, attn_analysis=attn_analysis)
+        k, v = self._apply_mea_head_mixing(k, v, attn_analysis=attn_analysis)
+        v    = self._apply_lucid_preconditioner(k, v, attention_mask, attn_analysis=attn_analysis)
 
         # Capture v_ref for XSA after MEA mixing and LUCID preconditioning.
         # This is the vector that actually participated in SDPA aggregation.
         v_ref = v if self.use_xsa else None
 
         # ── Affine-Scaled Attention ───────────────────────────────────────
+        # Active whenever use_affine_scaled_attention=True, regardless of
+        # attention backend.  Two code paths — same math, different execution:
+        #   eager : full weight access, attn_weights_pre/post_affine captured.
+        #   flash/sdpa: α·flash_out + β·V_cumsum, no weight tensors materialised.
         alpha = None
         beta  = None
-        use_affine = (
-            self.use_affine_scaled_attention
-            and self.config._attn_implementation == "eager"
-        )
+        use_affine = self.use_affine_scaled_attention
         if use_affine:
             alpha = linear_clipping(self.alpha_proj(hidden_states))   # [B, S, H]
             alpha = alpha.permute(0, 2, 1).unsqueeze(-1)              # [B, H, S, 1]
@@ -1476,14 +2010,31 @@ class NeoLLMAttention(nn.Module):
                         self.affine_momentum * self.alpha_ma
                         + (1.0 - self.affine_momentum) * batch_mean
                     )
+            if attn_analysis is not None:
+                attn_analysis.alpha_per_head   = alpha.detach()
+                attn_analysis.beta_per_head    = beta.detach()
+                attn_analysis.alpha_moving_avg = self.alpha_ma.detach()
 
         if use_affine:
-            attn_out, attn_weights = affine_scaled_eager_attention_forward(
-                self, q, k, v, attention_mask,
-                scaling=self.scaling, alpha=alpha, beta=beta,
-                dropout=0.0 if not self.training else self.attention_dropout,
-                **kwargs,
-            )
+            if self.config._attn_implementation == "eager":
+                # Eager: materialises softmax weights, full analysis available.
+                attn_out, attn_weights = affine_scaled_eager_attention_forward(
+                    self, q, k, v, attention_mask,
+                    scaling=self.scaling, alpha=alpha, beta=beta,
+                    dropout=0.0 if not self.training else self.attention_dropout,
+                    attn_analysis=attn_analysis,
+                    **kwargs,
+                )
+            else:
+                # Flash / SDPA: exact formula via V.cumsum, no weight tensors.
+                # attn_weights_pre/post_affine remain None in AnalysisState.
+                attn_out, attn_weights = affine_scaled_flash_attention_forward(
+                    self, q, k, v, attention_mask,
+                    scaling=self.scaling, alpha=alpha, beta=beta,
+                    dropout=0.0 if not self.training else self.attention_dropout,
+                    attn_analysis=attn_analysis,
+                    **kwargs,
+                )
         else:
             attn_fn = eager_attention_forward
             if self.config._attn_implementation != "eager":
@@ -1493,10 +2044,19 @@ class NeoLLMAttention(nn.Module):
                 dropout=0.0 if not self.training else self.attention_dropout,
                 scaling=self.scaling, **kwargs,
             )
+            if attn_analysis is not None:
+                attn_analysis.attn_weights = (
+                    attn_weights.detach() if attn_weights is not None else None
+                )
+
+        if attn_analysis is not None:
+            attn_analysis.attn_output_raw = attn_out.detach()
 
         attn_out = attn_out.reshape(*input_shape, -1, self.head_dim)
         if self.use_mea_attention:
             attn_out = self.mea_output_norm(attn_out)
+            if attn_analysis is not None:
+                attn_analysis.attn_output_post_mea_norm = attn_out.detach()
 
         # ── Exclusive Self Attention (position B, pre-routing) ────────────
         # Removes auto-position component before directional routing so that
@@ -1507,7 +2067,11 @@ class NeoLLMAttention(nn.Module):
             v_ref_t = v_ref_t.to(attn_out.dtype)
             proj    = (attn_out * v_ref_t).sum(dim=-1, keepdim=True)
             norm_sq = (v_ref_t * v_ref_t).sum(dim=-1, keepdim=True).clamp(min=self.xsa_eps)
-            attn_out = attn_out - (proj / norm_sq) * v_ref_t
+            xsa_comp = (proj / norm_sq) * v_ref_t
+            attn_out = attn_out - xsa_comp
+            if attn_analysis is not None:
+                attn_analysis.xsa_self_position_component = xsa_comp.detach()
+                attn_analysis.attn_output_post_xsa        = attn_out.detach()
 
         # ── Directional Routing (position C, post-XSA, pre-reshape) ──────
         # Suppresses cross-domain interference directions from the head output.
@@ -1515,14 +2079,26 @@ class NeoLLMAttention(nn.Module):
         # When use_xsa=False: directions span full head-space (no XSA pre-clean).
         # When use_directional_routing=False: this block is skipped entirely.
         if self.use_directional_routing:
-            attn_out = self._apply_directional_routing(attn_out, hidden_states)
+            attn_out = self._apply_directional_routing(
+                attn_out, hidden_states, attn_analysis=attn_analysis
+            )
 
         # ── Reshape → o_proj → Gated Attention gate → dropout ────────────
-        attn_out = attn_out.reshape(*input_shape, -1).contiguous()
-        attn_out = self.o_proj(attn_out * torch.sigmoid(gate))
-        attn_out = self.dropout(attn_out)
+        attn_out_flat = attn_out.reshape(*input_shape, -1).contiguous()
+        if attn_analysis is not None:
+            attn_analysis.attn_output_pre_gate = attn_out_flat.detach()
+            gate_sig = torch.sigmoid(gate)
+            attn_analysis.gate_sigmoid = gate_sig.detach()
+            attn_out_gated = self.o_proj(attn_out_flat * gate_sig)
+        else:
+            attn_out_gated = self.o_proj(attn_out_flat * torch.sigmoid(gate))
+
+        attn_out_gated = self.dropout(attn_out_gated)
 
-        return attn_out, attn_weights, current_layer_fan
+        if attn_analysis is not None:
+            attn_analysis.attn_output_final = attn_out_gated.detach()
+
+        return attn_out_gated, attn_weights, current_layer_fan
 class PolyNorm(nn.Module):
     def __init__(
         self,
@@ -1561,7 +2137,11 @@ class PolyNorm(nn.Module):
         out        = branch - alpha.to(branch.dtype) * proj_coeff * ref
         return self._norm(out)
 
-    def forward(self, x):
+    def forward(
+        self,
+        x: torch.Tensor,
+        analysis: Optional[PolyNormAnalysis] = None,
+    ) -> torch.Tensor:
         # Caché de potencias: x_sq reutilizado en x1 y x2; x_cu = x·x_sq evita pow(3)
         x_sq = x.pow(2)
         x_cu = x * x_sq
@@ -1571,9 +2151,20 @@ class PolyNorm(nn.Module):
         x2 = x_sq * (x_sq * x_sq).mean(-1, keepdim=True).add(self.eps).rsqrt()
         x3 = x_cu * (x_cu * x_cu).mean(-1, keepdim=True).add(self.eps).rsqrt()
 
+        if analysis is not None:
+            analysis.x1               = x1.detach()
+            analysis.x2_pre_exclusive = x2.detach()
+            analysis.x3_pre_exclusive = x3.detach()
+
         # Fuerzas exclusivas aprendibles
         alpha2, alpha3 = torch.sigmoid(self.exclusive_logits).unbind()
 
+        if analysis is not None:
+            analysis.alpha2  = alpha2.detach()
+            analysis.alpha3  = alpha3.detach()
+            analysis.weights = self.weight.detach()
+            analysis.bias    = self.bias.detach()
+
         # Precalcular ref (x1) en fp32 y su norma al cuadrado — compartido por x2 y x3
         x1_f        = x1.float()
         ref_norm_sq = x1_f.pow(2).sum(-1, keepdim=True).clamp_min(self.proj_eps)
@@ -1582,13 +2173,22 @@ class PolyNorm(nn.Module):
         x2 = self._exclusive(x2, x1, alpha2, x1_f, ref_norm_sq)
         x3 = self._exclusive(x3, x1, alpha3, x1_f, ref_norm_sq)
 
-        return (
+        if analysis is not None:
+            analysis.x2_post_exclusive = x2.detach()
+            analysis.x3_post_exclusive = x3.detach()
+
+        output = (
             self.weight[0] * x3
             + self.weight[1] * x2
             + self.weight[2] * x1
             + self.bias
         )
 
+        if analysis is not None:
+            analysis.output = output.detach()
+
+        return output
+
 
 class NeoLLMMLP(nn.Module):
     """MLP with FANformer integration and Learnable Multipliers."""
@@ -1614,11 +2214,34 @@ class NeoLLMMLP(nn.Module):
         self.act_fn  = PolyNorm(exclusive_init=0.05)
         self.dropout = nn.Dropout(config.dropout_rate)
 
-    def forward(self, x):
-        x_fan = self.fan_layer(x)
-        return self.down_proj(
-            self.dropout(self.act_fn(self.gate_proj(x_fan)) * self.up_proj(x_fan))
-        )
+    def forward(
+        self,
+        x: torch.Tensor,
+        analysis: Optional[MLPAnalysis] = None,
+    ) -> torch.Tensor:
+        fan_a = analysis.fan if analysis is not None else None
+        x_fan = self.fan_layer(x, analysis=fan_a)
+
+        gate_out = self.gate_proj(x_fan)
+        up_out   = self.up_proj(x_fan)
+
+        if analysis is not None:
+            analysis.gate_proj_output = gate_out.detach()
+            analysis.up_proj_output   = up_out.detach()
+
+        poly_a   = analysis.polynorm if analysis is not None else None
+        act_out  = self.act_fn(gate_out, analysis=poly_a)
+        act_x_up = act_out * up_out
+
+        if analysis is not None:
+            analysis.act_times_up = act_x_up.detach()
+
+        result = self.down_proj(self.dropout(act_x_up))
+
+        if analysis is not None:
+            analysis.output = result.detach()
+
+        return result
 
 
 class NeoLLMDecoderLayer(GradientCheckpointingLayer):
@@ -1681,6 +2304,8 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         sources: list,
         partial: torch.Tensor,
         query: torch.Tensor,
+        analysis_slot: Optional[AttnResAnalysis] = None,
+        analysis_key: Optional[str] = None,
     ) -> torch.Tensor:
         """
         Depth-wise softmax attention over preceding layer outputs.
@@ -1722,9 +2347,14 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         B_vals: Optional[torch.Tensor] = None,
         attn_res_sources: Optional[list] = None,
         attn_res_partial: Optional[torch.Tensor] = None,
+        layer_analysis: Optional[LayerAnalysis] = None,
         output_attentions: Optional[bool] = False,
         **kwargs: Unpack[FlashAttentionKwargs],
     ) -> Tuple:
+        # ── Snapshot input ────────────────────────────────────────────────
+        if layer_analysis is not None:
+            layer_analysis.hidden_states_input = hidden_states.detach()
+
         # ── Attention Residuals: compute pre-attention input ──────────────
         # When active, the input to the attention sublayer is no longer the
         # raw hidden_states (accumulated residual) but a softmax-weighted
@@ -1732,41 +2362,68 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         # attn_res_partial carries the intra-block standard residual that
         # connects the attention and MLP sublayers within this layer.
         # When inactive, flow is identical to the original.
+        ar_analysis = layer_analysis.attn_res if layer_analysis is not None else None
         if self.use_attn_res and attn_res_sources is not None and attn_res_partial is not None:
-            h_attn = self._attn_res(attn_res_sources, attn_res_partial,
-                                     self.attn_res_query_attn)
+            h_attn = self._attn_res(
+                attn_res_sources, attn_res_partial, self.attn_res_query_attn,
+                ar_analysis, "attn",
+            )
             residual_attn = attn_res_partial
         else:
             h_attn = hidden_states
             residual_attn = hidden_states
 
         # ── Attention block ───────────────────────────────────────────────
-        h_attn = self.lns_attn(self.input_layernorm(h_attn))
+        sn_pre = layer_analysis.seednorm_pre_attn if layer_analysis is not None else None
+        h_normed = self.input_layernorm(h_attn, analysis=sn_pre)
+        h_lns    = self.lns_attn(h_normed)
+        if layer_analysis is not None:
+            layer_analysis.lns_attn_output = h_lns.detach()
 
         hidden_states, attn_weights, self.current_layer_fan = self.self_attn(
-            hidden_states=h_attn,
+            hidden_states=h_lns,
             attention_mask=attention_mask,
             position_embeddings=position_embeddings,
             first_layer_fan=first_layer_fan,
+            attn_analysis=layer_analysis.attention if layer_analysis is not None else None,
             **kwargs,
         )
-        h_tilde = self.gpas_attn(residual_attn + hidden_states)
+
+        if layer_analysis is not None:
+            layer_analysis.attn_contribution = hidden_states.detach()
+
+        gpas_attn_a = layer_analysis.gpas_attn if layer_analysis is not None else None
+        h_tilde = self.gpas_attn(residual_attn + hidden_states, analysis=gpas_attn_a)
+
+        if layer_analysis is not None:
+            layer_analysis.h_tilde = h_tilde.detach()
 
         # ── Attention Residuals: compute pre-MLP input ────────────────────
         # After attention, the partial sum is updated with h_tilde.
         # The pre-MLP AttnRes attends over the same sources but with h_tilde
         # as the current partial — capturing the within-layer attention output.
         if self.use_attn_res and attn_res_sources is not None:
-            h_mlp = self._attn_res(attn_res_sources, h_tilde,
-                                    self.attn_res_query_mlp)
+            h_mlp = self._attn_res(
+                attn_res_sources, h_tilde, self.attn_res_query_mlp,
+                ar_analysis, "mlp",
+            )
             residual_mlp = h_tilde
         else:
             h_mlp = h_tilde
             residual_mlp = h_tilde
 
         # ── MLP block ─────────────────────────────────────────────────────
-        h_mlp    = self.lns_mlp(self.post_attention_layernorm(h_mlp))
-        delta_m  = self.mlp(h_mlp)
+        sn_post = layer_analysis.seednorm_post_attn if layer_analysis is not None else None
+        h_normed2 = self.post_attention_layernorm(h_mlp, analysis=sn_post)
+        h_lns2    = self.lns_mlp(h_normed2)
+        if layer_analysis is not None:
+            layer_analysis.lns_mlp_output = h_lns2.detach()
+
+        mlp_a   = layer_analysis.mlp if layer_analysis is not None else None
+        delta_m = self.mlp(h_lns2, analysis=mlp_a)
+
+        if layer_analysis is not None:
+            layer_analysis.mlp_contribution = delta_m.detach()
 
         # ── JTok-M injection (additive alongside MLP residual) ────────────
         aux_stats = None
@@ -1776,12 +2433,18 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
             z_flat     = z_tilde.reshape(-1, z_tilde.shape[-1])
             B_flat     = B_vals.reshape(-1, B_vals.shape[-2], B_vals.shape[-1])
 
-            delta_r, aux_stats = self.jtokm(h_flat, z_flat, B_flat)
+            jtokm_a = layer_analysis.jtokm if layer_analysis is not None else None
+            delta_r, aux_stats = self.jtokm(h_flat, z_flat, B_flat, analysis=jtokm_a)
             delta_r = delta_r.reshape(orig_shape)
 
-            hidden_states = self.gpas_mlp(residual_mlp + delta_m + delta_r)
+            gpas_mlp_a    = layer_analysis.gpas_mlp if layer_analysis is not None else None
+            hidden_states = self.gpas_mlp(residual_mlp + delta_m + delta_r, analysis=gpas_mlp_a)
         else:
-            hidden_states = self.gpas_mlp(residual_mlp + delta_m)
+            gpas_mlp_a    = layer_analysis.gpas_mlp if layer_analysis is not None else None
+            hidden_states = self.gpas_mlp(residual_mlp + delta_m, analysis=gpas_mlp_a)
+
+        if layer_analysis is not None:
+            layer_analysis.hidden_states_output = hidden_states.detach()
 
         outputs = (hidden_states,)
         if output_attentions:
@@ -1968,6 +2631,30 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
         else:
             self.embed_tokens = value
 
+    def _build_layer_analysis(self) -> LayerAnalysis:
+        """
+        Construct a LayerAnalysis with sub-objects pre-allocated for every
+        component that is active in the current config.
+
+        Fields that correspond to disabled flags remain None — the analysis
+        consumer can check for None without needing to inspect config flags.
+        Called once per layer per forward when analysis is active.
+        """
+        cfg = self.config
+        return LayerAnalysis(
+            seednorm_pre_attn  = SeeDNormAnalysis(),
+            seednorm_post_attn = SeeDNormAnalysis(),
+            attention          = AttentionAnalysis(fan=FANAnalysis()),
+            mlp                = MLPAnalysis(
+                fan      = FANAnalysis(),
+                polynorm = PolyNormAnalysis(),
+            ),
+            gpas_attn = GPASAnalysis(),
+            gpas_mlp  = GPASAnalysis(),
+            jtokm     = JTokMAnalysis() if cfg.use_jtokm else None,
+            attn_res  = AttnResAnalysis() if getattr(cfg, "use_attn_res", False) else None,
+        )
+
     def forward(
         self,
         input_ids: Optional[torch.LongTensor] = None,
@@ -1977,6 +2664,7 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
         output_hidden_states: Optional[bool] = None,
         output_attentions: Optional[bool] = None,
         return_dict: Optional[bool] = None,
+        analysis_state: Optional[AnalysisState] = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> Tuple:
         output_hidden_states = (
@@ -1996,12 +2684,18 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
         z_tilde = None
         B_vals  = None
 
+        gen_a = (
+            analysis_state.generator
+            if analysis_state is not None and self.config.use_token_generator
+            else None
+        )
+
         if inputs_embeds is None:
             if self.config.use_token_generator:
                 if self.config.use_jtokm:
                     # Return internals for reuse by JTok-M surfaces
                     inputs_embeds, z_tilde, B_vals = self.token_generator(
-                        input_ids, return_internals=True
+                        input_ids, return_internals=True, analysis=gen_a
                     )
                     # Reshape to [batch, seq, d_seed] and [batch, seq, d_seed, n_knots]
                     z_tilde = z_tilde.reshape(*input_ids.shape, self.config.generator_d_seed)
@@ -2011,10 +2705,13 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
                         self.config.generator_num_knots,
                     )
                 else:
-                    inputs_embeds = self.token_generator(input_ids)
+                    inputs_embeds = self.token_generator(input_ids, analysis=gen_a)
             else:
                 inputs_embeds = self.embed_tokens(input_ids)
 
+        if analysis_state is not None:
+            analysis_state.embeddings = inputs_embeds.detach()
+
         if position_ids is None:
             position_ids = torch.arange(
                 0, inputs_embeds.shape[1], device=inputs_embeds.device
@@ -2061,6 +2758,10 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
                 else 1   # Full AttnRes: every layer is its own "block"
             )
 
+        # Pre-allocate per-layer analysis list when analysis is active
+        if analysis_state is not None:
+            analysis_state.layers = []
+
         for layer_idx, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]):
             if output_hidden_states:
                 all_hidden_states = all_hidden_states + (hidden_states,)
@@ -2077,6 +2778,13 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
                 attn_res_sources = attn_res_sources + [attn_res_partial]
                 attn_res_partial = hidden_states  # start new block from current output
 
+            # Build per-layer analysis container (only in eval + analysis mode)
+            layer_analysis = None
+            if analysis_state is not None:
+                layer_analysis           = self._build_layer_analysis()
+                layer_analysis.layer_idx = layer_idx
+                analysis_state.layers.append(layer_analysis)
+
             layer_outputs = decoder_layer(
                 hidden_states,
                 position_embeddings=position_embeddings,
@@ -2086,6 +2794,7 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
                 B_vals=B_vals,
                 attn_res_sources=attn_res_sources,
                 attn_res_partial=attn_res_partial if use_attn_res else None,
+                layer_analysis=layer_analysis,
                 output_attentions=output_attentions,
                 **kwargs,
             )
@@ -2111,6 +2820,14 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
         if output_hidden_states:
             all_hidden_states = all_hidden_states + (hidden_states,)
 
+        # ── Finalise analysis snapshot ─────────────────────────────────────
+        if analysis_state is not None:
+            analysis_state.final_hidden_states     = hidden_states.detach()
+            analysis_state.jtokm_aux_stats         = all_aux_stats if self.config.use_jtokm else None
+            analysis_state.attn_res_sources_final  = (
+                attn_res_sources if use_attn_res else None
+            )
+
         if not return_dict:
             return tuple(
                 v for v in [hidden_states, None, all_hidden_states, all_attentions]
@@ -2153,6 +2870,30 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
         total_loss = CE_loss + L_aux
 
     where L_aux = λ · n_e · (1/L) · Σ_ℓ Σ_i p_i^ℓ · f_i^ℓ
+
+    ── Analysis mode ──────────────────────────────────────────────────────────
+    Analysis is NEVER active during training (model.training=True).
+    It is opt-in at inference time:
+
+        model.eval()
+        model.enable_analysis()        # arm the system
+        with torch.no_grad():
+            _ = model(input_ids)
+        state = model.last_analysis    # AnalysisState with all internals
+
+        model.disable_analysis()       # disarm — zero overhead again
+        # or: model.train()            # training always disarms automatically
+
+    last_analysis is replaced at each forward call when analysis is armed.
+    last_analysis is None between enable_analysis() and the first forward,
+    and cleared by disable_analysis().
+
+    Example field access:
+        state.layers[3].attention.alpha_per_head      # Affine-Scaled α
+        state.layers[0].mlp.polynorm.x2_post_exclusive
+        state.generator.z_tilde                       # Leviathan latent
+        state.layers[2].jtokm.routing_weights         # JTok-M router
+        state.layers[5].attn_res.weights_pre_attn     # AttnRes softmax
     """
 
     _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
@@ -2166,8 +2907,70 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
         if config.use_token_generator:
             self._tied_weights_keys = {}
 
+        # ── Analysis infrastructure ───────────────────────────────────────
+        # _analysis_armed: set by enable_analysis() / disable_analysis().
+        # last_analysis:   populated after each forward when armed + eval.
+        # Neither is an nn.Parameter or buffer — zero effect on training.
+        self._analysis_armed: bool               = False
+        self.last_analysis:   Optional[AnalysisState] = None
+
         self.post_init()
 
+    # ── Public analysis API ───────────────────────────────────────────────
+
+    def enable_analysis(self) -> None:
+        """
+        Arm the analysis system.
+
+        Has no effect during training — analysis is always off when
+        model.training=True, regardless of this flag. Safe to call at any
+        point without affecting the training loop.
+
+        Usage::
+
+            model.eval()
+            model.enable_analysis()
+            with torch.no_grad():
+                _ = model(input_ids)
+            state = model.last_analysis
+        """
+        self._analysis_armed = True
+
+    def disable_analysis(self) -> None:
+        """
+        Disarm the analysis system and clear last_analysis.
+
+        After this call, forward passes produce zero analysis overhead.
+        """
+        self._analysis_armed = False
+        self.last_analysis   = None
+
+    def _make_analysis_state(
+        self,
+        input_ids: Optional[torch.Tensor],
+    ) -> Optional[AnalysisState]:
+        """
+        Decide whether to build an AnalysisState for this forward pass.
+
+        Returns None (zero cost) when:
+          - analysis is not armed, OR
+          - the model is in training mode  (self.training=True).
+        Otherwise returns a fresh AnalysisState with top-level optional
+        fields pre-allocated according to the active config flags.
+        """
+        if not self._analysis_armed or self.training:
+            return None
+        cfg   = self.config
+        return AnalysisState(
+            input_ids              = input_ids.detach() if input_ids is not None else None,
+            generator              = GeneratorAnalysis() if cfg.use_token_generator else None,
+            layers                 = None,   # filled by NeoLLMModel.forward
+            jtokm_aux_stats        = [] if cfg.use_jtokm else None,
+            attn_res_sources_final = [] if getattr(cfg, "use_attn_res", False) else None,
+        )
+
+    # ── Standard model API ────────────────────────────────────────────────
+
     def get_input_embeddings(self):
         return self.model.get_input_embeddings()
 
@@ -2217,6 +3020,9 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
         return_dict: Optional[bool] = None,
         **kwargs: Unpack[TransformersKwargs],
     ) -> CausalLMOutputWithPast:
+        # ── Build analysis container (None during training or when disarmed) ──
+        analysis_state = self._make_analysis_state(input_ids)
+
         model_out = self.model(
             input_ids=input_ids,
             attention_mask=attention_mask,
@@ -2224,6 +3030,7 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
             inputs_embeds=inputs_embeds,
             output_hidden_states=output_hidden_states,
             return_dict=return_dict,
+            analysis_state=analysis_state,
             **kwargs,
         )
 
@@ -2262,6 +3069,11 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
             )
             logits = self.lm_head(hidden_states[:, slice_indices, :])
 
+        # ── Finalise and store analysis state ─────────────────────────────
+        if analysis_state is not None:
+            analysis_state.logits = logits.detach() if logits is not None else None
+            self.last_analysis    = analysis_state
+
         return CausalLMOutputWithPast(
             loss=loss,
             logits=logits,
@@ -2286,6 +3098,18 @@ __all__ = [
     "VectorMultiplier",
     "LinearWithMultipliers",
     "MEAHeadSeeDNorm",
+    # Analysis dataclasses — exported so external tools can type-hint against them
+    "AnalysisState",
+    "LayerAnalysis",
+    "AttentionAnalysis",
+    "MLPAnalysis",
+    "FANAnalysis",
+    "SeeDNormAnalysis",
+    "GPASAnalysis",
+    "PolyNormAnalysis",
+    "JTokMAnalysis",
+    "AttnResAnalysis",
+    "GeneratorAnalysis",
 ]
 
 AutoConfig.register("neollm", NeoLLMConfig)