Update modeling_neollm.py

modeling_neollm.py

@@ -239,6 +239,10 @@ class NeoLLMBaseModelOutputWithPast(BaseModelOutputWithPast):
     jtokm_aux_stats: Optional[
         Tuple[Tuple[torch.Tensor, torch.Tensor, torch.Tensor], ...]
     ] = None
+    # StackMemory diagnostics packed as [num_layers, 11].  Column 0 retains
+    # the differentiable action entropy used by the auxiliary StackTrans loss;
+    # the remaining columns are detached monitoring reductions.
+    stack_metrics: Optional[torch.Tensor] = None
 
 logger = logging.get_logger(__name__)
 _NEOLLM_FA_USE_TOP_LEFT_MASK = flash_attn_supports_top_left_mask()
@@ -403,6 +407,394 @@ class EmbeddingWithMultipliers(nn.Module):
 # ==================== LEVIATHAN CONTINUOUS TOKEN GENERATOR ====================
 
 
+# Sampled dynamics: mathematical definitions and interpretation
+# ----------------------------------------------------------
+# Enable through USE_DYNAMICS_METRICS in train.py before model construction.
+# train.py supplies config.dynamics_metrics_enabled and dynamics_sample_tokens;
+# a standalone caller may set those same attributes before constructing a model.
+# Collection requires native CUDA/Triton checkpoints, training mode and enabled
+# autograd. It is detached, read-only, and never contributes to the objective.
+#
+# Sampling/aggregation contract:
+# Flatten batch and sequence to N rows, choose S=min(N, sample_tokens) positions
+# r_j=floor(j*N/S), then exclude invalid sampled rows. There is no resampling
+# to replace masked rows. Explicit attention validity excludes padding; MEAP
+# replacement rows are excluded from generator diagnostics. Token IDs alone
+# are not a padding test (EOS may share the padding ID). JToK uses its existing
+# validity mask. Let V be the remaining sampled rows and A(f)=mean_{r in V} f_r.
+# Every row statistic below is reduced with A, except sample_valid_rows=|V|.
+# Thus A(RMS(v_r)) is NOT RMS over all flattened activations. Empty V produces
+# neutral zeros; consult sample_valid_rows before interpreting any metric.
+# Deterministic position sampling is not an unbiased dataset estimate and can
+# alias positional structure. The cache describes the last gradient-enabled
+# microbatch/recomputation, not an optimizer-step, epoch or validation average.
+# Arithmetic uses FP32 on detached, actually stored native tensors. Nonfinite
+# values are not silently repaired; interpret RMS/angles with finite fractions.
+#
+# Leviathan metrics (prefix dynamics/leviathan/):
+# Let e_r be the generated embedding, z_r the seed, and m_r the concatenation
+# of stored CP product modes across heads/ranks. Define RMS(v)=sqrt(mean(v^2)),
+# Z(v)=mean(1[v=0]), NF(v)=mean(1[not finite(v)]).
+#   embedding_rms = A(RMS(e)); tracks generator output scale.
+#   seed_rms = A(RMS(z)); tracks the shared seed representation's scale.
+#   product_mode_rms = A(RMS(m)); tracks the CP branch's numerical magnitude.
+#   embedding_nonfinite_fraction = A(NF(e)); flags invalid generated outputs.
+#   product_mode_zero_fraction = A(Z(m)); flags stored zero/underflow/collapse,
+#       without distinguishing their causes or the separate residual branch.
+#   product_mode_nonfinite_fraction = A(NF(m)); localizes invalid CP products.
+#   coordinate_saturation_fraction = A(mean_{head,d} 1[c<=.01 or c>=.99]),
+#       c=sigmoid(.5*(gamma*xhat+beta)). xhat/gamma/beta are native checkpoints;
+#       c is reconstructed only for sampled rows in FP32, before any activation
+#       cast. This is a boundary/conditioning proxy, not the exact stored
+#       coordinate's threshold test or proof that spline derivatives vanish.
+# Seed, mode and embedding scales separate representation growth from product
+# attenuation and projection effects. Their trends alone do not prove learning.
+#
+# JToK versus JToK-M: keep their distinct forward semantics when comparing.
+# x is the incoming MLP update delta_m, not the full decoder hidden state;
+# y is the update returned by this module, d=y-x, s is the learned scaler.
+# With n(v)=v/(||v||_2+norm_eps), on valid rows:
+#   JToK:   y=x*(1+s*n(S(z_tilde)))       [elementwise modulation].
+#   JToK-M: y=x+a*s*n(sum_k p_k*S_{e_k}(z_tilde)) [additive residual],
+#            a=1/sqrt(2L), e_k=TopK(router_logits),
+#            p_k=sigmoid(logit_{e_k})/max(sum_j sigmoid(logit_{e_j}),norm_eps).
+# Native stored output rounding is part of the observed d. The observer does
+# not reconstruct an ideal gate/residual or change normalization/routing.
+#
+# Shared metrics (prefix dynamics/jtok|jtokm/layer_<index>/), eps=1e-12:
+#   input_rms=A(RMS(x)), output_rms=A(RMS(y)), update_rms=A(RMS(d));
+#       separate incoming scale, outgoing scale and effective intervention.
+#   relative_update_norm=A(||d||_2/max(||x||_2,eps)); unbounded effect size,
+#       sensitive to near-zero input. Read alongside zero_input_fraction.
+#   usage_index=A(||d||_2/max(||x||_2+||d||_2,eps)); bounded [0,1] for finite
+#       inputs, zero for an identity intervention. It measures visible effect,
+#       NOT probability of use, benefit, gradient flow or semantic relevance.
+#   input_output_cosine=A(clip(<x,y>/max(||x||_2*||y||_2,eps),-1,1));
+#       detects direction preservation/reversal separately from scale changes.
+#   update_alignment=A(clip(<x,d>/max(||x||_2*||d||_2,eps),-1,1));
+#       positive/negative values indicate reinforcement/cancellation of x.
+#   orthogonal_update_fraction=A(1-alignment_r^2 if ||d||_2>0 else 0);
+#       for nonzero, well-scaled x this is the fraction of d's energy
+#       orthogonal to x. With x=0,d!=0 the convention is 1 (no reference
+#       direction); epsilon-dominated angles are proxies, not exact geometry.
+#       It is the mean of per-row fractions, NOT 1-A(alignment)^2.
+#   changed_element_fraction=A(mean_h 1[y_h!=x_h]); visible dtype-level
+#       intervention, not an FP32-master update or usefulness test.
+#   zero_input_fraction=A(1[||x||_2=0]); qualifies norm ratios and angles.
+#   output_nonfinite_fraction=A(NF(y)); output numerical-health check.
+#   product_mode_zero_fraction=A(Z(m)), product_mode_nonfinite_fraction=A(NF(m));
+#       m contains all stored selected-expert CP modes (one expert for JToK).
+#   coordinate_saturation_fraction=A(mean_d 1[z_tilde_d<=.01 or >=.99]);
+#       observes the supplied coordinate, not the generator head coordinates.
+# Jointly usage, alignment and orthogonal fraction distinguish a small/identity
+# intervention, reinforcement/cancellation and direction-changing corrections.
+# They cannot establish why a token needs that correction; causal usefulness
+# requires a controlled ablation and task/loss evidence, not these proxies.
+#
+# Additional JToK-M routing metrics (not emitted for plain JToK):
+# For selected stored weights p, b=sum_k p_k and pi_k=p_k/b when b>0 (else 0):
+#   selected_weight_entropy=A(-sum_k pi_k*log(pi_k)); entropy in nats over the
+#       selected K, not all E experts. K=1 is necessarily zero. Diagnostic pi
+#       renormalization does not modify the actual p used by the model.
+#   selected_max_weight=A(max_k p_k); concentration of actual mixture weights.
+#   selected_weight_sum=A(b), zero_selected_mass_fraction=A(1[b=0]); reveal
+#       epsilon-limited or underflowed mixture mass; entropy alone can hide it.
+# Let f_e=sum_{r in V,k}1[e_k(r)=e]/(|V|*K), selection counts, NOT weight mass:
+#   expert_<e>_selection_fraction=f_e; sampled expert traffic.
+#   load_cv=std_population(f)/max(mean(f),eps); imbalance across experts.
+#   load_entropy_normalized=-sum_e f_e*log(max(f_e,1e-30))/max(log(E),eps);
+#       coverage/uniformity on [0,1] when V is nonempty and E>1; E=1 yields 0.
+#   active_experts=sum_e 1[f_e>0]; observed coverage, not causal contribution.
+# These characterize routing collapse/specialization candidates, not expert
+# quality. Low load entropy may be intentional specialization. Always compare
+# with usage, update geometry, numerical health and train/validation loss.
+#
+# Cost/API contract: each observed forward samples at most sample_tokens rows,
+# reduces on GPU and caches only compact detached FP32 scalars, never token IDs
+# or activation vectors. There are extra GPU launches/workspaces, so cheap does
+# not mean free. get_dynamics_metrics() exposes device scalars outside compiled
+# forward; train.py batches their host transfer only when logging. Evaluation
+# does not collect these training snapshots or relabel them as eval metrics.
+try:
+    import triton
+    import triton.language as tl
+except (ImportError, ModuleNotFoundError):
+    triton = tl = None
+
+LEV_NAMES = (
+    "sample_valid_rows", "embedding_rms", "embedding_nonfinite_fraction",
+    "seed_rms", "coordinate_saturation_fraction", "product_mode_rms",
+    "product_mode_zero_fraction", "product_mode_nonfinite_fraction",
+)
+JTOK_NAMES = (
+    "sample_valid_rows", "input_rms", "output_rms", "update_rms",
+    "relative_update_norm", "input_output_cosine", "changed_element_fraction",
+    "product_mode_zero_fraction", "product_mode_nonfinite_fraction",
+    "coordinate_saturation_fraction", "selected_weight_entropy",
+    "selected_max_weight", "output_nonfinite_fraction", "zero_input_fraction",
+    "usage_index", "update_alignment", "orthogonal_update_fraction",
+    "selected_weight_sum", "zero_selected_mass_fraction",
+)
+
+
+if triton is not None:
+    @triton.jit
+    def _dynamics_reduce_kernel(rows, output, S: tl.constexpr, C: tl.constexpr,
+                                BS: tl.constexpr):
+        col = tl.program_id(0)
+        r = tl.arange(0, BS)
+        valid = tl.load(rows + r * C, r < S, 0)
+        count = tl.sum(valid, 0)
+        values = tl.load(rows + r * C + col, r < S, 0)
+        value = tl.sum(tl.where(valid > 0, values, 0), 0) / tl.maximum(count, 1)
+        tl.store(output + col, tl.where(col == 0, count, value))
+
+    @triton.jit
+    def _jtok_dynamics_rows_kernel(delta, out, z, modes, experts, weights, valid,
+                                   rows, N: tl.constexpr, S: tl.constexpr,
+                                   H: tl.constexpr, D: tl.constexpr, M: tl.constexpr,
+                                   K: tl.constexpr, E: tl.constexpr, BH: tl.constexpr,
+                                   BD: tl.constexpr, BM: tl.constexpr,
+                                   BK: tl.constexpr, HAS_MASK: tl.constexpr):
+        sample = tl.program_id(0)
+        row = sample * N // S
+        is_valid = tl.full((), True, tl.int1)
+        if HAS_MASK:
+            is_valid = tl.load(valid + row)
+        h = tl.arange(0, BH)
+        x = tl.load(delta + row * H + h, h < H, 0).to(tl.float32)
+        y = tl.load(out + row * H + h, h < H, 0).to(tl.float32)
+        diff = y - x
+        xn = tl.sqrt(tl.sum(x * x, 0))
+        yn = tl.sqrt(tl.sum(y * y, 0))
+        dn = tl.sqrt(tl.sum(diff * diff, 0))
+        cosine = tl.minimum(1., tl.maximum(-1., tl.sum(x * y, 0) / tl.maximum(xn * yn, 1e-12)))
+        d = tl.arange(0, BD)
+        coord = tl.load(z + row * D + d, d < D, 0).to(tl.float32)
+        saturation = tl.sum(((coord <= .01) | (coord >= .99)) & (d < D), 0) / D
+        m = tl.arange(0, BM)
+        mode = tl.load(modes + row * K * M + m, m < K * M, 0).to(tl.float32)
+        mz = tl.sum((mode == 0) & (m < K * M), 0) / (K * M)
+        mn = tl.sum(((mode != mode) | (tl.abs(mode) == float('inf'))) & (m < K * M), 0) / (K * M)
+        k = tl.arange(0, BK)
+        p = tl.load(weights + row * K + k, k < K, 0).to(tl.float32)
+        mass = tl.sum(p, 0)
+        probability = p / tl.where(mass > 0, mass, 1.)
+        entropy = -tl.sum(tl.where(probability > 0,
+                                   probability * tl.log(tl.maximum(probability, 1e-30)), 0), 0)
+        max_weight = tl.max(p, 0)
+        nonfinite = tl.sum(((y != y) | (tl.abs(y) == float('inf'))) & (h < H), 0) / H
+        changed = tl.sum((y != x) & (h < H), 0) / H
+        alignment = tl.minimum(1., tl.maximum(-1., tl.sum(x * diff, 0) / tl.maximum(xn * dn, 1e-12)))
+        orthogonal = tl.where(dn > 0, tl.maximum(0., 1. - alignment * alignment), 0.)
+        C: tl.constexpr = 19 + E
+        base = rows + sample * C
+        tl.store(base, is_valid.to(tl.float32))
+        tl.store(base + 1, xn / tl.sqrt(float(H)))
+        tl.store(base + 2, yn / tl.sqrt(float(H)))
+        tl.store(base + 3, dn / tl.sqrt(float(H)))
+        tl.store(base + 4, dn / tl.maximum(xn, 1e-12))
+        tl.store(base + 5, cosine)
+        tl.store(base + 6, changed)
+        tl.store(base + 7, mz)
+        tl.store(base + 8, mn)
+        tl.store(base + 9, saturation)
+        tl.store(base + 10, entropy)
+        tl.store(base + 11, max_weight)
+        tl.store(base + 12, nonfinite)
+        tl.store(base + 13, (xn == 0).to(tl.float32))
+        tl.store(base + 14, dn / tl.maximum(xn + dn, 1e-12))
+        tl.store(base + 15, alignment)
+        tl.store(base + 16, orthogonal)
+        tl.store(base + 17, mass)
+        tl.store(base + 18, (mass == 0).to(tl.float32))
+        route = tl.load(experts + row * K + k, k < K, -1)
+        for expert in range(E):
+            fraction = tl.sum((route == expert) & (k < K), 0) / K
+            tl.store(base + 19 + expert, fraction)
+
+    @triton.jit
+    def _lev_dynamics_rows_kernel(embedding, seed, xhat, norm_weight, norm_bias,
+                                  modes, valid, rows, N: tl.constexpr, S: tl.constexpr,
+                                  H: tl.constexpr, D: tl.constexpr, HEADS: tl.constexpr,
+                                  R: tl.constexpr, BH: tl.constexpr, BD: tl.constexpr,
+                                  BM: tl.constexpr, HAS_MASK: tl.constexpr):
+        sample = tl.program_id(0)
+        row = sample * N // S
+        is_valid = tl.full((), True, tl.int1)
+        if HAS_MASK:
+            is_valid = tl.load(valid + row)
+        h = tl.arange(0, BH)
+        value = tl.load(embedding + row * H + h, h < H, 0).to(tl.float32)
+        erms = tl.sqrt(tl.sum(value * value, 0) / H)
+        en = tl.sum(((value != value) | (tl.abs(value) == float('inf'))) & (h < H), 0) / H
+        d = tl.arange(0, BD)
+        z = tl.load(seed + row * D + d, d < D, 0).to(tl.float32)
+        zrms = tl.sqrt(tl.sum(z * z, 0) / D)
+        saturated = tl.full((), 0., tl.float32)
+        for head in range(HEADS):
+            x = tl.load(xhat + (head * N + row) * D + d, d < D, 0).to(tl.float32)
+            w = tl.load(norm_weight + head * D + d, d < D, 0).to(tl.float32)
+            b = tl.load(norm_bias + head * D + d, d < D, 0).to(tl.float32)
+            coordinate = tl.sigmoid(.5 * (x * w + b))
+            saturated += tl.sum(((coordinate <= .01) | (coordinate >= .99)) & (d < D), 0)
+        m = tl.arange(0, BM)
+        mode = tl.load(modes + row * HEADS * R + m, m < HEADS * R, 0).to(tl.float32)
+        mrms = tl.sqrt(tl.sum(mode * mode, 0) / (HEADS * R))
+        mz = tl.sum((mode == 0) & (m < HEADS * R), 0) / (HEADS * R)
+        mn = tl.sum(((mode != mode) | (tl.abs(mode) == float('inf'))) & (m < HEADS * R), 0) / (HEADS * R)
+        base = rows + sample * 8
+        tl.store(base, is_valid.to(tl.float32))
+        tl.store(base + 1, erms)
+        tl.store(base + 2, en)
+        tl.store(base + 3, zrms)
+        tl.store(base + 4, saturated / (HEADS * D))
+        tl.store(base + 5, mrms)
+        tl.store(base + 6, mz)
+        tl.store(base + 7, mn)
+
+
+def _check_sampling(max_samples: int) -> None:
+    validate_dynamics_samples(max_samples)
+    if max_samples == 0:
+        raise ValueError("diagnostics max_samples must be in [1, 4096]")
+
+
+def validate_dynamics_samples(samples: int) -> None:
+    """Zero disables collection; reject invalid budgets before native work."""
+    if type(samples) is not int or not 0 <= samples <= 4096:
+        raise ValueError("dynamics_samples must be an integer in [0, 4096]")
+
+
+def _check_tensors(primary: torch.Tensor, *others: torch.Tensor) -> None:
+    for value in (primary, *others):
+        if value.device != primary.device or not value.is_contiguous():
+            raise ValueError("diagnostics tensors must be contiguous on the same device")
+        if value.dtype not in (torch.float16, torch.bfloat16, torch.float32):
+            raise TypeError("diagnostics floating tensors require FP16, BF16 or FP32")
+
+
+def _check_mask(valid: torch.Tensor, n: int, device: torch.device) -> None:
+    if (valid.dtype != torch.bool or valid.device != device or valid.ndim != 1
+            or valid.numel() not in (0, n) or not valid.is_contiguous()):
+        raise ValueError("diagnostics mask must be a contiguous bool vector of length zero or N on the input device")
+
+
+def _reduce(rows: torch.Tensor, columns: int) -> torch.Tensor:
+    output = torch.empty(columns, device=rows.device, dtype=torch.float32)
+    _dynamics_reduce_kernel[(columns,)](
+        rows, output, rows.shape[0], columns,
+        triton.next_power_of_2(rows.shape[0]), num_warps=4)
+    return output
+
+
+@torch.library.custom_op("neollm_dynamics::jtok_dynamics", mutates_args=())
+def jtok_dynamics(delta: torch.Tensor, output: torch.Tensor, z: torch.Tensor,
+                  modes: torch.Tensor, experts: torch.Tensor, weights: torch.Tensor,
+                  valid: torch.Tensor, num_experts: int, max_samples: int) -> torch.Tensor:
+    _check_sampling(max_samples)
+    if type(num_experts) is not int or not 1 <= num_experts <= 1024:
+        raise ValueError("diagnostics num_experts must be an integer in [1, 1024]")
+    if delta.ndim != 2 or delta.shape[1] == 0 or output.shape != delta.shape:
+        raise ValueError("diagnostics input/output must have shape [N, H] with H > 0")
+    n, hidden = delta.shape
+    if z.ndim != 2 or z.shape[0] != n or z.shape[1] == 0:
+        raise ValueError("diagnostics coordinates must have shape [N, D] with D > 0")
+    if (modes.ndim != 3 or modes.shape[0] != n or modes.shape[2] == 0
+            or not 1 <= modes.shape[1] <= num_experts):
+        raise ValueError("diagnostics modes must have shape [N, K, M] with 1 <= K <= E and M > 0")
+    if experts.shape != modes.shape[:2] or weights.shape != experts.shape:
+        raise ValueError("diagnostics routes/weights must have shape [N, K]")
+    if (experts.dtype not in (torch.int32, torch.int64)
+            or experts.device != delta.device or not experts.is_contiguous()):
+        raise ValueError("diagnostics routes must be contiguous integer tensors on the input device")
+    _check_tensors(delta, output, z, modes, weights)
+    _check_mask(valid, n, delta.device)
+    if triton is None or not delta.is_cuda:
+        raise RuntimeError("native JToK diagnostics require CUDA/Triton")
+    samples = min(n, max_samples)
+    columns = len(JTOK_NAMES) + num_experts
+    if not samples:
+        return torch.zeros(columns, device=delta.device, dtype=torch.float32)
+    rows = torch.empty((samples, columns), device=delta.device, dtype=torch.float32)
+    _jtok_dynamics_rows_kernel[(samples,)](
+        delta, output, z, modes, experts, weights, valid, rows, n, samples,
+        hidden, z.shape[1], modes.shape[2], modes.shape[1], num_experts,
+        triton.next_power_of_2(hidden), triton.next_power_of_2(z.shape[1]),
+        triton.next_power_of_2(modes.shape[1] * modes.shape[2]),
+        triton.next_power_of_2(modes.shape[1]), valid.numel() != 0,
+        num_warps=4)
+    return _reduce(rows, columns)
+
+
+@jtok_dynamics.register_fake
+def _jtok_dynamics_fake(delta, output, z, modes, experts, weights, valid,
+                        num_experts, max_samples):
+    return torch.empty(len(JTOK_NAMES) + num_experts, device=delta.device, dtype=torch.float32)
+
+
+@torch.library.custom_op("neollm_dynamics::leviathan_dynamics", mutates_args=())
+def leviathan_dynamics(embedding: torch.Tensor, seed: torch.Tensor, xhat: torch.Tensor,
+                       norm_weight: torch.Tensor, norm_bias: torch.Tensor,
+                       modes: torch.Tensor, valid: torch.Tensor, max_samples: int) -> torch.Tensor:
+    _check_sampling(max_samples)
+    if embedding.ndim < 2 or embedding.shape[-1] == 0:
+        raise ValueError("diagnostics embedding must have shape [..., H] with H > 0")
+    n, hidden = embedding.numel() // embedding.shape[-1], embedding.shape[-1]
+    if (xhat.ndim != 3 or xhat.shape[0] == 0 or xhat.shape[1] != n
+            or xhat.shape[2] == 0):
+        raise ValueError("diagnostics xhat must have shape [heads, N, D] with positive heads and D")
+    heads, _, d = xhat.shape
+    if seed.shape != (n, d) or norm_weight.shape != (heads, d) or norm_bias.shape != (heads, d):
+        raise ValueError("diagnostics seed/norm metadata do not match xhat")
+    if modes.ndim != 3 or modes.shape[:2] != (n, heads) or modes.shape[2] == 0:
+        raise ValueError("diagnostics modes must have shape [N, heads, rank] with rank > 0")
+    _check_tensors(embedding, seed, xhat, norm_weight, norm_bias, modes)
+    _check_mask(valid, n, embedding.device)
+    if triton is None or not embedding.is_cuda:
+        raise RuntimeError("native Leviathan diagnostics require CUDA/Triton")
+    samples = min(n, max_samples)
+    if not samples:
+        return torch.zeros(len(LEV_NAMES), device=embedding.device, dtype=torch.float32)
+    rows = torch.empty((samples, len(LEV_NAMES)), device=embedding.device, dtype=torch.float32)
+    rank = modes.shape[-1]
+    _lev_dynamics_rows_kernel[(samples,)](
+        embedding, seed, xhat, norm_weight, norm_bias, modes, valid, rows,
+        n, samples, hidden, d, heads, rank, triton.next_power_of_2(hidden),
+        triton.next_power_of_2(d), triton.next_power_of_2(heads * rank),
+        valid.numel() != 0, num_warps=4)
+    return _reduce(rows, len(LEV_NAMES))
+
+
+@leviathan_dynamics.register_fake
+def _leviathan_dynamics_fake(embedding, seed, xhat, norm_weight, norm_bias,
+                             modes, valid, max_samples):
+    return torch.empty(len(LEV_NAMES), device=embedding.device, dtype=torch.float32)
+
+
+def unpack_jtok_dynamics(packed: torch.Tensor, num_experts: int, *, mixture: bool):
+    """Expose device scalars; the caller chooses when to transfer to its logger."""
+    result = dict(zip(JTOK_NAMES, packed[:len(JTOK_NAMES)].unbind()))
+    if not mixture:
+        result.pop("selected_weight_entropy")
+        result.pop("selected_max_weight")
+        result.pop("selected_weight_sum")
+        result.pop("zero_selected_mass_fraction")
+        return result
+    fractions = packed[len(JTOK_NAMES):]
+    mean = fractions.mean()
+    result["load_cv"] = fractions.std(unbiased=False) / mean.clamp_min(1e-12)
+    result["load_entropy_normalized"] = (
+        -(fractions * fractions.clamp_min(1e-30).log()).sum()
+        / torch.tensor(float(num_experts), device=packed.device).log().clamp_min(1e-12))
+    result["active_experts"] = (fractions > 0).sum().float()
+    for expert, value in enumerate(fractions.unbind()):
+        result[f"expert_{expert}_selection_fraction"] = value
+    return result
+
+
+
+
 class LeviathanGenerator(nn.Module):
     """
     Learned Embedding Vectorization (LEV) layer - Leviathan input embedding
@@ -451,6 +843,11 @@ class LeviathanGenerator(nn.Module):
         # is unavailable; callers can also disable the optional route per
         # instance without changing the model/state-dict contract.
         self.use_leviathan_triton = bool(_LEV_KERNEL_AVAILABLE)
+        self.dynamics_samples = (
+            int(getattr(config, "dynamics_sample_tokens", 64))
+            if bool(getattr(config, "dynamics_metrics_enabled", False)) else 0
+        )
+        self._last_dynamics_metrics = None
         vocab_size = config.vocab_size
         hidden_size = config.hidden_size
         self.d_seed = config.generator_d_seed
@@ -639,6 +1036,18 @@ class LeviathanGenerator(nn.Module):
         prod_sign = 1.0 - 2.0 * (num_neg % 2).float()  # [N, krank]
         return prod_sign * torch.exp(log_mag)  # [N, krank]
 
+    def _record_dynamics(self, token_ids, embedding, checkpoints, valid_mask, mask_token_id):
+        seed, xhat, modes = checkpoints
+        diag_mask = (token_ids.new_empty((0,), dtype=torch.bool) if valid_mask is None
+                     else valid_mask.reshape(-1).contiguous())
+        if mask_token_id is not None:
+            unmasked = token_ids.reshape(-1).ne(int(mask_token_id))
+            diag_mask = unmasked if diag_mask.numel() == 0 else diag_mask & unmasked
+        packed = leviathan_dynamics(
+            embedding.detach(), seed, xhat, self.head_norm_weight.detach(),
+            self.head_norm_bias.detach(), modes, diag_mask, self.dynamics_samples)
+        self._last_dynamics_metrics = packed.detach().clone()
+
     # Keep this method traceable. The compiler-safe LEV entry point registers
     # an opaque custom op tagged cudagraph_unsafe: Dynamo keeps the LEV call
     # in the compiled graph while Inductor excludes only the unsafe Triton
@@ -651,6 +1060,7 @@ class LeviathanGenerator(nn.Module):
         meap_mask_embedding: Optional[torch.Tensor] = None,
         meap_mask_token_id: Optional[int] = None,
         return_jtok_geometry: bool = False,
+        dynamics_valid_mask: Optional[torch.Tensor] = None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
         """
         Generate embeddings from discrete token indices.
@@ -665,6 +1075,11 @@ class LeviathanGenerator(nn.Module):
             where ``z_tilde`` is the shared differentiable Leviathan/JTok
             coordinate with shape ``[*token_ids.shape, d_seed]``.
         """
+        collect_dynamics = bool(self.dynamics_samples and self.training and torch.is_grad_enabled())
+        dynamics_kwargs = (
+            {"return_dynamics_checkpoints": True}
+            if collect_dynamics else {}
+        )
         has_meap = (
             meap_mask_embedding is not None or meap_mask_token_id is not None
         )
@@ -704,7 +1119,7 @@ class LeviathanGenerator(nn.Module):
         )
         require_leviathan_triton = bool(
             getattr(self.config, "_require_leviathan_triton", False)
-        )
+        ) or collect_dynamics
         if require_leviathan_triton and not leviathan_kernel_ready:
             raise RuntimeError(
                 "This run requires the Triton Leviathan kernel, but the "
@@ -749,24 +1164,38 @@ class LeviathanGenerator(nn.Module):
                             "but the installed CCE package lacks the geometry "
                             "adapter."
                         )
-                    embedding, seed = _leviathan_embedding_with_seed_compiler_safe(
+                    result = _leviathan_embedding_with_seed_compiler_safe(
                         token_ids,
                         params,
                         self.config,
                         self.knot_grid,
                         mask_embedding=meap_mask_embedding,
                         mask_token_id=meap_mask_token_id,
+                        **dynamics_kwargs,
                     )
+                    if collect_dynamics:
+                        embedding, seed, checkpoints = result
+                        self._record_dynamics(token_ids, embedding, checkpoints,
+                                              dynamics_valid_mask, meap_mask_token_id)
+                    else:
+                        embedding, seed = result
                     embedding = embedding.reshape(*token_ids.shape, self.hidden_size)
                     return embedding, self._jtok_coordinate_from_seed(seed)
-                return _leviathan_embedding_compiler_safe(
+                result = _leviathan_embedding_compiler_safe(
                     token_ids,
                     params,
                     self.config,
                     self.knot_grid,
                     mask_embedding=meap_mask_embedding,
                     mask_token_id=meap_mask_token_id,
+                    **dynamics_kwargs,
                 )
+                if collect_dynamics:
+                    embedding, checkpoints = result
+                    self._record_dynamics(token_ids, embedding, checkpoints,
+                                          dynamics_valid_mask, meap_mask_token_id)
+                    return embedding
+                return result
             except (TypeError, ValueError, AttributeError):
                 # Only validation/configuration failures use the reference
                 # implementation. A CUDA launch error must propagate: doing
@@ -870,6 +1299,11 @@ class LeviathanJTok(nn.Module):
         super().__init__()
         self.layer_idx = int(layer_idx)
         self.use_mixture = bool(getattr(config, "use_jtokm", False))
+        self.dynamics_samples = (
+            int(getattr(config, "dynamics_sample_tokens", 64))
+            if bool(getattr(config, "dynamics_metrics_enabled", False)) else 0
+        )
+        self._last_dynamics_metrics = None
         self.jtok_kernel_backend = str(
             getattr(config, "jtok_kernel_backend", "torch")
         ).strip().lower()
@@ -877,6 +1311,8 @@ class LeviathanJTok(nn.Module):
             raise ValueError(
                 "`jtok_kernel_backend` must be either `torch` or `triton`."
             )
+        if self.dynamics_samples and self.jtok_kernel_backend != "triton":
+            raise RuntimeError("Native dynamics metrics require the strict Triton JTok backend")
         if self.jtok_kernel_backend == "triton" and not _JTOK_KERNEL_AVAILABLE:
             raise RuntimeError(
                 "`jtok_kernel_backend='triton'` was requested, but the "
@@ -1104,6 +1540,7 @@ class LeviathanJTok(nn.Module):
         Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]],
     ]:
         """Apply JTok/JTok-M and optionally return training-only router stats."""
+        collect_dynamics = bool(self.dynamics_samples and self.training and torch.is_grad_enabled())
         if self.jtok_kernel_backend == "triton":
             # Dispatch only: this branch does not change compiler mode,
             # CUDA-graph policy, training state, optimizer behavior, or any
@@ -1116,7 +1553,7 @@ class LeviathanJTok(nn.Module):
                     "JTok Triton adapter became unavailable after module "
                     "construction; refusing an implicit Torch fallback."
                 )
-            return _apply_neollm_jtok(
+            result = _apply_neollm_jtok(
                 self,
                 delta_m,
                 z_tilde,
@@ -1124,7 +1561,22 @@ class LeviathanJTok(nn.Module):
                 valid_mask=valid_mask,
                 compute_aux=bool(compute_aux and self.use_mixture),
                 backend="triton",
-            )
+                **({"return_dynamics_checkpoints": True} if collect_dynamics else {}),
+            )
+            if collect_dynamics:
+                output, aux, checkpoints = result
+                modes, routes, weights = checkpoints
+                n = delta_m.numel() // self.hidden_size
+                mask = (delta_m.new_empty((0,), dtype=torch.bool) if valid_mask is None
+                        else valid_mask.reshape(-1).to(dtype=torch.bool).contiguous())
+                packed = jtok_dynamics(
+                    delta_m.detach().reshape(n, self.hidden_size).contiguous(),
+                    output.detach().reshape(n, self.hidden_size).contiguous(),
+                    z_tilde.detach().reshape(n, self.d_seed).contiguous(), modes, routes, weights,
+                    mask, self.num_experts, self.dynamics_samples)
+                self._last_dynamics_metrics = packed.detach().clone()
+                return output, aux
+            return result
 
         orig_shape = delta_m.shape
         N = delta_m.numel() // self.hidden_size
@@ -1286,7 +1738,11 @@ class StackMemory(nn.Module):
         self.down_proj = nn.Linear(config.hidden_size, config.stack_d_model)
         self.up_proj = nn.Linear(config.stack_d_model, config.hidden_size)
         self.action_head = nn.Linear(config.stack_d_model, 3 * self.num_mem_heads)
-        self.gate_proj = nn.Linear(self.stack_dim, 1)
+        # A scalar bias is exactly softmax-invariant because the same value is
+        # added to every stack slot.  Keeping it would create a parameter tensor
+        # with an analytically zero gradient, which is especially undesirable for
+        # per-tensor GradientStabilizer scaling.
+        self.gate_proj = nn.Linear(self.stack_dim, 1, bias=False)
         self.res_weight = nn.Parameter(torch.ones(1))
 
         self.cache_size = getattr(config, "stack_memory_cache_size", 2048)
@@ -1311,24 +1767,37 @@ class StackMemory(nn.Module):
             [mask[:, :, :, 1:], torch.zeros_like(mask[:, :, :, :1])], dim=3
         )
 
-        action_weights = actions.unsqueeze(-1).unsqueeze(-1)
+        # Keep stack values in the model dtype (normally BF16), but keep the
+        # differentiable control mask in FP32.  The latter is important because
+        # the global-read mask bias can be extremely large (e.g. -1e9).
+        action_weights_stack = actions.to(stack.dtype).unsqueeze(-1).unsqueeze(-1)
+        action_weights_mask = actions.to(mask.dtype).unsqueeze(-1)
         stacks = torch.stack([push_stack, pop_stack, stack], dim=3)
         masks = torch.stack([push_mask, pop_mask, mask], dim=3)
 
-        new_stack = (stacks * action_weights).sum(dim=3)
-        new_mask = (masks * action_weights.squeeze(-1)).sum(dim=3)
+        new_stack = (stacks * action_weights_stack).sum(dim=3)
+        new_mask = (masks * action_weights_mask).sum(dim=3)
         return new_stack, new_mask
 
-    def forward(self, hidden_states, stack, mask):
+    def forward(self, hidden_states, stack, mask, token_mask=None):
         batch_size, seq_len, _ = hidden_states.shape
 
+        # The stack occupancy is a differentiable control state.  Keep it in
+        # FP32 across layers/callers even when model activations are BF16.
+        mask = mask.float()
+
         # OLMo stack_v3 value/action path: project the complete hidden state into
         # the total low-rank stack width, then split that representation by head.
         new_hidden_states = self.down_proj(hidden_states)
 
         action_logits = self.action_head(new_hidden_states) / math.sqrt(self.stack_dim)
+        # Stack actions are control probabilities.  Evaluate softmax in FP32 so
+        # the recurrent mask is never quantized to BF16 before later layers use
+        # it.  Gradients still flow normally into the BF16 action_head weights.
         actions = F.softmax(
-            action_logits.reshape(batch_size, seq_len, self.num_mem_heads, 3),
+            action_logits.float().reshape(
+                batch_size, seq_len, self.num_mem_heads, 3
+            ),
             dim=-1,
         )
 
@@ -1338,21 +1807,99 @@ class StackMemory(nn.Module):
 
         new_stack, new_mask = self._vectorized_update(stack, mask, actions, k_values)
 
-        gate_scores = self.gate_proj(new_stack).squeeze(-1)
-        # Keep NeoLLM's requested softened source mask exactly at -80.
-        gate_weights = F.softmax(gate_scores + (1 - new_mask) * -80.0, dim=-1)
+        # The global-read control path is evaluated explicitly in FP32.  This
+        # preserves the original differentiable graph (no detach) while avoiding
+        # BF16 quantization of both the soft mask and the very large mask bias.
+        stack_fp32 = new_stack.float()
+        gate_scores = F.linear(
+            stack_fp32,
+            self.gate_proj.weight.float(),
+            bias=None,
+        ).squeeze(-1)
+        gate_weights = F.softmax(
+            gate_scores + (1.0 - new_mask) * -320,
+            dim=-1,
+        )
 
-        memory_output = (new_stack * gate_weights.unsqueeze(-1)).sum(dim=3)
+        memory_output = (stack_fp32 * gate_weights.unsqueeze(-1)).sum(dim=3)
+        memory_output = memory_output.to(new_hidden_states.dtype)
         memory_output = memory_output.reshape(batch_size, seq_len, -1)
         memory_output = self.up_proj(memory_output)
 
         output = memory_output * self.res_weight + hidden_states
 
+        # Compact StackTrans statistics.  Entropy is evaluated in FP32 for
+        # numerical stability; only its scalar mean keeps autograd because it is
+        # used by the auxiliary objective.  All other diagnostics are detached
+        # before they leave this layer so monitoring does not retain extra graphs.
+        actions_fp32 = actions.float()
+        action_entropy_map = -(
+            actions_fp32 * torch.log(actions_fp32.clamp_min(1e-12))
+        ).sum(dim=-1)  # [B,S,H]
+        action_max_map = actions_fp32.max(dim=-1).values
+
+        gate_fp32 = gate_weights.float()
+        gate_entropy_map = -(
+            gate_fp32 * torch.log(gate_fp32.clamp_min(1e-12))
+        ).sum(dim=-1)  # [B,S,H]
+        gate_max_map = gate_fp32.max(dim=-1).values
+        gate_effective_slots_map = gate_entropy_map.exp()
+
+        mask_fp32 = new_mask.float()
+        expected_depth_map = mask_fp32.sum(dim=-1)  # [B,S,H]
+        mask_fill_map = expected_depth_map / float(self.stack_slots)
+        top_occupancy_map = mask_fp32[..., 0]
+
+        if token_mask is None:
+            action_entropy_mean = action_entropy_map.mean()
+            action_probs_mean = actions_fp32.mean(dim=(0, 1, 2))
+            action_max_mean = action_max_map.mean()
+            expected_depth_mean = expected_depth_map.mean()
+            mask_fill_mean = mask_fill_map.mean()
+            top_occupancy_mean = top_occupancy_map.mean()
+            gate_entropy_mean = gate_entropy_map.mean()
+            gate_max_mean = gate_max_map.mean()
+            gate_effective_slots_mean = gate_effective_slots_map.mean()
+        else:
+            valid = token_mask.to(
+                device=hidden_states.device, dtype=torch.float32
+            ).unsqueeze(-1)  # [B,S,1]
+            denom = (valid.sum() * float(self.num_mem_heads)).clamp_min(1.0)
+            action_entropy_mean = (action_entropy_map * valid).sum() / denom
+            action_probs_mean = (
+                actions_fp32 * valid.unsqueeze(-1)
+            ).sum(dim=(0, 1, 2)) / denom
+            action_max_mean = (action_max_map * valid).sum() / denom
+            expected_depth_mean = (expected_depth_map * valid).sum() / denom
+            mask_fill_mean = (mask_fill_map * valid).sum() / denom
+            top_occupancy_mean = (top_occupancy_map * valid).sum() / denom
+            gate_entropy_mean = (gate_entropy_map * valid).sum() / denom
+            gate_max_mean = (gate_max_map * valid).sum() / denom
+            gate_effective_slots_mean = (
+                gate_effective_slots_map * valid
+            ).sum() / denom
+
+        stack_metrics = torch.stack(
+            (
+                action_entropy_mean,
+                action_probs_mean[0].detach(),
+                action_probs_mean[1].detach(),
+                action_probs_mean[2].detach(),
+                action_max_mean.detach(),
+                expected_depth_mean.detach(),
+                mask_fill_mean.detach(),
+                top_occupancy_mean.detach(),
+                gate_entropy_mean.detach(),
+                gate_max_mean.detach(),
+                gate_effective_slots_mean.detach(),
+            )
+        )
+
         if self.training and self.enable_cache:
             self._update_cache(k_values.detach(), actions.detach())
 
         # Preserve the complete token-local stack for the next decoder layer.
-        return output, new_stack, new_mask
+        return output, new_stack, new_mask, stack_metrics
 
     def _update_cache(self, k_values, actions):
         seq_len = k_values.shape[1]
@@ -1376,7 +1923,7 @@ class StackMemory(nn.Module):
         if mask.ndim == 3:
             mask = mask.unsqueeze(1)
 
-        output, new_stack, new_mask = self.forward(
+        output, new_stack, new_mask, _stack_metrics = self.forward(
             hidden_state.unsqueeze(1), stack, mask
         )
         return output.squeeze(1), new_stack.squeeze(1), new_mask.squeeze(1)
@@ -5487,7 +6034,13 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         hidden_states: torch.Tensor,
         stack_memory: Optional[torch.Tensor],
         stack_memory_mask: Optional[torch.Tensor],
-    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
+        token_mask: Optional[torch.Tensor] = None,
+    ) -> Tuple[
+        torch.Tensor,
+        Optional[torch.Tensor],
+        Optional[torch.Tensor],
+        Optional[torch.Tensor],
+    ]:
         """
         Apply this layer's differentiable StackMemory module before attention.
 
@@ -5497,13 +6050,15 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         attention implementation remains unchanged.
         """
         if (not self.use_stack_memory) or self.stack_memory is None:
-            return hidden_states, stack_memory, stack_memory_mask
+            return hidden_states, stack_memory, stack_memory_mask, None
         if stack_memory is None or stack_memory_mask is None:
             raise ValueError(
                 "StackMemory is enabled, but stack_memory/stack_memory_mask "
                 "were not initialized by NeoLLMModel.forward."
             )
-        return self.stack_memory(hidden_states, stack_memory, stack_memory_mask)
+        return self.stack_memory(
+            hidden_states, stack_memory, stack_memory_mask, token_mask=token_mask
+        )
 
     def _attn_res(
         self,
@@ -5624,6 +6179,7 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         position_embeddings: tuple[torch.Tensor, torch.Tensor],
         stack_memory: Optional[torch.Tensor] = None,
         stack_memory_mask: Optional[torch.Tensor] = None,
+        stack_token_mask: Optional[torch.Tensor] = None,
         attention_mask: Optional[torch.Tensor] = None,
         first_layer_fan: Optional[torch.Tensor] = None,
         output_attentions: Optional[bool] = False,
@@ -5647,9 +6203,15 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
         if self.siamese_attn_input_norm is not None:
             h_attn = self.siamese_attn_input_norm(h_attn)
 
+        stack_metrics = None
         if self.use_stack_memory:
-            h_attn, stack_memory, stack_memory_mask = self.apply_stack_memory(
-                h_attn, stack_memory, stack_memory_mask
+            h_attn, stack_memory, stack_memory_mask, stack_metrics = (
+                self.apply_stack_memory(
+                    h_attn,
+                    stack_memory,
+                    stack_memory_mask,
+                    token_mask=stack_token_mask,
+                )
             )
         jtok_router_state = h_attn
 
@@ -5706,7 +6268,7 @@ class NeoLLMDecoderLayer(GradientCheckpointingLayer):
 
         outputs = (x_next, y_next)
         if self.use_stack_memory:
-            outputs += (stack_memory, stack_memory_mask)
+            outputs += (stack_memory, stack_memory_mask, stack_metrics)
         if jtok_aux_stats is not None:
             outputs += jtok_aux_stats
         if output_attentions:
@@ -6637,6 +7199,22 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
             )
         jtok_z_tilde = None
 
+        # Diagnostic validity is separate from attention/loss/JTok masks.
+        # In particular, EOS may also be the padding ID; never infer padding
+        # from IDs. Exclude MEAP replacements even if SBE owns the epilogue.
+        dynamics_kwargs = {}
+        if (bool(getattr(self.config, "dynamics_metrics_enabled", False))
+                and self.training and torch.is_grad_enabled() and input_ids is not None):
+            if attention_mask is not None:
+                if attention_mask.ndim != 2 or attention_mask.shape != input_ids.shape:
+                    raise ValueError("Dynamics collection requires an explicit 2D attention validity mask")
+                dynamics_mask = attention_mask.to(dtype=torch.bool)
+            else:
+                dynamics_mask = torch.ones_like(input_ids, dtype=torch.bool)
+            if meap_active:
+                dynamics_mask = dynamics_mask & input_ids.ne(int(self.config.meap_mask_token_id))
+            dynamics_kwargs = {"dynamics_valid_mask": dynamics_mask}
+
         # ── Embedding stage ────────────────────────────────────────────────
         meap_mask_embedding = None
         meap_mask_token_id = None
@@ -6665,6 +7243,7 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
                             meap_mask_embedding=meap_mask_embedding,
                             meap_mask_token_id=meap_mask_token_id,
                             return_jtok_geometry=True,
+                            **dynamics_kwargs,
                         )
                     else:
                         # Keep the pre-JTok call signature for custom or
@@ -6673,15 +7252,17 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
                             input_ids,
                             meap_mask_embedding=meap_mask_embedding,
                             meap_mask_token_id=meap_mask_token_id,
+                            **dynamics_kwargs,
                         )
                 else:
                     if use_jtok:
                         generator_result = self.token_generator(
                             input_ids,
                             return_jtok_geometry=True,
+                            **dynamics_kwargs,
                         )
                     else:
-                        generator_result = self.token_generator(input_ids)
+                        generator_result = self.token_generator(input_ids, **dynamics_kwargs)
                 if use_jtok:
                     inputs_embeds, jtok_z_tilde = generator_result
                 else:
@@ -6820,12 +7401,17 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
                 .to(dtype=hidden_states.dtype)
                 .expand(batch_size, seq_len, -1, -1, -1)
             )
-            stack_memory_mask = self.memory_mask.detach().expand(
-                batch_size, seq_len, -1, -1
+            # Keep the differentiable StackMemory occupancy/control state in
+            # FP32 even when the live model parameters/activations are BF16.
+            stack_memory_mask = (
+                self.memory_mask.detach()
+                .to(device=hidden_states.device, dtype=torch.float32)
+                .expand(batch_size, seq_len, -1, -1)
             )
         else:
             stack_memory = None
             stack_memory_mask = None
+        all_stack_metrics = () if use_stack_memory else None
 
         position_embeddings = self.rotary_emb(hidden_states, position_ids)
         self.first_layer_fan = (
@@ -6901,11 +7487,15 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
                 attn_res_partial = hidden_states  # start new block from current output
 
             if use_stack_memory and not use_siamesenorm:
-                hidden_states, stack_memory, stack_memory_mask = (
+                hidden_states, stack_memory, stack_memory_mask, stack_layer_metrics = (
                     decoder_layer.apply_stack_memory(
-                        hidden_states, stack_memory, stack_memory_mask
+                        hidden_states,
+                        stack_memory,
+                        stack_memory_mask,
+                        token_mask=attention_mask,
                     )
                 )
+                all_stack_metrics = all_stack_metrics + (stack_layer_metrics,)
 
             if use_siamesenorm:
                 layer_outputs = decoder_layer.forward_siamesenorm(
@@ -6914,6 +7504,7 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
                     position_embeddings=position_embeddings,
                     stack_memory=stack_memory,
                     stack_memory_mask=stack_memory_mask,
+                    stack_token_mask=attention_mask,
                     attention_mask=causal_mask,
                     first_layer_fan=self.first_layer_fan,
                     output_attentions=output_attentions,
@@ -6929,7 +7520,9 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
                 if use_stack_memory:
                     stack_memory = layer_outputs[2]
                     stack_memory_mask = layer_outputs[3]
-                    extras_start = 4
+                    stack_layer_metrics = layer_outputs[4]
+                    all_stack_metrics = all_stack_metrics + (stack_layer_metrics,)
+                    extras_start = 5
                 else:
                     extras_start = 2
             else:
@@ -6990,12 +7583,20 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
         if output_hidden_states:
             all_hidden_states = all_hidden_states + (hidden_states,)
 
+        stack_metrics = (
+            torch.stack(all_stack_metrics, dim=0)
+            if all_stack_metrics is not None and len(all_stack_metrics) > 0
+            else None
+        )
+
         if not return_dict:
             base_outputs = tuple(
                 v
                 for v in [hidden_states, None, all_hidden_states, all_attentions]
                 if v is not None
             )
+            if stack_metrics is not None:
+                base_outputs = base_outputs + (stack_metrics,)
             if all_jtokm_aux_stats is not None:
                 base_outputs = base_outputs + (all_jtokm_aux_stats,)
             if output_tweo_activations:
@@ -7008,6 +7609,7 @@ class NeoLLMModel(NeoLLMPreTrainedModel):
             hidden_states=all_hidden_states,
             attentions=all_attentions,
             jtokm_aux_stats=all_jtokm_aux_stats,
+            stack_metrics=stack_metrics,
         )
         if output_tweo_activations:
             return outputs, tweo_activations
@@ -8840,6 +9442,20 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
         self._last_nextlat_mse_loss = None
         self._last_nextlat_kl_loss = None
         self._last_jtokm_aux_loss = None
+        self._last_stack_entropy_loss = None
+        self._last_stack_action_entropy = None
+        self._last_stack_action_entropy_normalized = None
+        self._last_stack_push_prob = None
+        self._last_stack_pop_prob = None
+        self._last_stack_noop_prob = None
+        self._last_stack_action_max_prob = None
+        self._last_stack_expected_depth = None
+        self._last_stack_mask_fill_fraction = None
+        self._last_stack_top_occupancy = None
+        self._last_stack_gate_entropy = None
+        self._last_stack_gate_entropy_normalized = None
+        self._last_stack_gate_max_prob = None
+        self._last_stack_gate_effective_slots = None
         self._last_total_loss = None
 
         if config.use_token_generator:
@@ -8850,6 +9466,50 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
     def get_input_embeddings(self):
         return self.model.get_input_embeddings()
 
+    def get_dynamics_metrics(self):
+        """Detached device scalars from the last gradient-enabled training forward.
+
+        Read outside compiled forward at logging cadence. Sampling describes
+        the last microbatch/recomputation, not an average over optimizer steps.
+        This method does not synchronize, transfer activations, or change loss.
+        """
+        if not bool(getattr(self.config, "dynamics_metrics_enabled", False)):
+            return {}
+        metrics = {}
+        generator = getattr(self.model, "token_generator", None)
+        packed = getattr(generator, "_last_dynamics_metrics", None)
+        if packed is not None:
+            metrics.update({f"dynamics/leviathan/{name}": value.detach()
+                            for name, value in zip(LEV_NAMES, packed.unbind())})
+        for layer in self.model.layers:
+            module = getattr(layer, "jtok", None)
+            packed = getattr(module, "_last_dynamics_metrics", None)
+            if packed is None:
+                continue
+            kind = "jtokm" if module.use_mixture else "jtok"
+            prefix = f"dynamics/{kind}/layer_{module.layer_idx}/"
+            metrics.update({prefix + name: value.detach() for name, value in
+                            unpack_jtok_dynamics(packed, module.num_experts,
+                                                 mixture=module.use_mixture).items()})
+        return metrics
+
+    def get_dynamics_parameter_groups(self):
+        """Only investigated generator/bridge/surface/router parameters, no backbone."""
+        groups = {}
+        generator = getattr(self.model, "token_generator", None)
+        if generator is not None:
+            for name, parameter in generator.named_parameters():
+                if parameter.requires_grad:
+                    groups[f"leviathan/{name}"] = (parameter,)
+        for layer in self.model.layers:
+            module = getattr(layer, "jtok", None)
+            if module is not None:
+                kind = "jtokm" if module.use_mixture else "jtok"
+                for name, parameter in module.named_parameters():
+                    if parameter.requires_grad:
+                        groups[f"{kind}/layer_{module.layer_idx}/{name}"] = (parameter,)
+        return groups
+
     def set_input_embeddings(self, value):
         self.model.set_input_embeddings(value)
 
@@ -9025,6 +9685,7 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
         tweo_activations = None
         hidden_states_tuple = None
         jtokm_aux_stats = None
+        stack_metrics = None
         if isinstance(model_out, tuple):
             outputs = model_out[0]
             if tweo_enabled and len(model_out) >= 2:
@@ -9041,6 +9702,13 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                     if tuple_candidates:
                         hidden_states_tuple = tuple_candidates[0]
                 for candidate in model_out[1:]:
+                    if (
+                        bool(getattr(self.config, "use_stack_memory", False))
+                        and isinstance(candidate, torch.Tensor)
+                        and candidate.ndim == 2
+                        and candidate.shape[-1] == 11
+                    ):
+                        stack_metrics = candidate
                     if (
                         isinstance(candidate, tuple)
                         and candidate
@@ -9048,7 +9716,6 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                         and len(candidate[0]) == 3
                     ):
                         jtokm_aux_stats = candidate
-                        break
             elif isinstance(outputs, tuple):
                 hidden_states = outputs[0]
                 if len(outputs) > 2:
@@ -9057,11 +9724,13 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                 hidden_states = outputs.last_hidden_state
                 hidden_states_tuple = outputs.hidden_states
                 jtokm_aux_stats = getattr(outputs, "jtokm_aux_stats", None)
+                stack_metrics = getattr(outputs, "stack_metrics", None)
         else:
             outputs = model_out
             hidden_states = outputs.last_hidden_state
             hidden_states_tuple = outputs.hidden_states
             jtokm_aux_stats = getattr(outputs, "jtokm_aux_stats", None)
+            stack_metrics = getattr(outputs, "stack_metrics", None)
 
         loss = None
         ntp_loss = None
@@ -9085,6 +9754,7 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
         nextlat_mse_loss = None
         nextlat_kl_loss = None
         jtokm_aux_loss = None
+        stack_entropy_loss = None
         self._last_ntp_loss = None
         self._last_ntp_ce_unweighted = None
         self._last_mile_reweighting_delta = None
@@ -9111,6 +9781,20 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
         self._last_nextlat_mse_loss = None
         self._last_nextlat_kl_loss = None
         self._last_jtokm_aux_loss = None
+        self._last_stack_entropy_loss = None
+        self._last_stack_action_entropy = None
+        self._last_stack_action_entropy_normalized = None
+        self._last_stack_push_prob = None
+        self._last_stack_pop_prob = None
+        self._last_stack_noop_prob = None
+        self._last_stack_action_max_prob = None
+        self._last_stack_expected_depth = None
+        self._last_stack_mask_fill_fraction = None
+        self._last_stack_top_occupancy = None
+        self._last_stack_gate_entropy = None
+        self._last_stack_gate_entropy_normalized = None
+        self._last_stack_gate_max_prob = None
+        self._last_stack_gate_effective_slots = None
         self._last_total_loss = None
         if labels is not None:
             if self.ntp_loss_backend == "liger":
@@ -9166,6 +9850,24 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
                     ntp_loss = ntp_result
             loss = ntp_loss
 
+            # StackTrans entropy regularizer.  The backbone returns one compact
+            # metric row per decoder layer; column 0 is the differentiable mean
+            # entropy H(push,pop,noop).  Averaging across layers keeps lambda
+            # invariant to model depth, batch size, sequence length, and head count.
+            if bool(getattr(self.config, "use_stack_memory", False)):
+                if stack_metrics is None or stack_metrics.numel() == 0:
+                    raise RuntimeError(
+                        "StackMemory is active but no stack metrics were returned."
+                    )
+                stack_action_entropy = stack_metrics[:, 0].mean()
+                stack_entropy_weight = torch.as_tensor(
+                    float(getattr(self.config, "stack_entropy_loss_weight", 1e-3)),
+                    device=stack_action_entropy.device,
+                    dtype=torch.float32,
+                )
+                stack_entropy_loss = stack_entropy_weight * stack_action_entropy
+                loss = loss + stack_entropy_loss
+
             if tweo_enabled:
                 if not tweo_activations:
                     raise ValueError(
@@ -9451,6 +10153,30 @@ class NeoLLMForCausalLM(NeoLLMPreTrainedModel, GenerationMixin):
             self._last_jtokm_aux_loss = (
                 jtokm_aux_loss.detach() if jtokm_aux_loss is not None else None
             )
+            if bool(getattr(self.config, "use_stack_memory", False)):
+                detached_stack = stack_metrics.detach().float()
+                action_entropy = detached_stack[:, 0].mean()
+                gate_entropy = detached_stack[:, 8].mean()
+                self._last_stack_entropy_loss = (
+                    stack_entropy_loss.detach()
+                    if stack_entropy_loss is not None
+                    else action_entropy.new_zeros(())
+                )
+                self._last_stack_action_entropy = action_entropy
+                self._last_stack_action_entropy_normalized = action_entropy / math.log(3.0)
+                self._last_stack_push_prob = detached_stack[:, 1].mean()
+                self._last_stack_pop_prob = detached_stack[:, 2].mean()
+                self._last_stack_noop_prob = detached_stack[:, 3].mean()
+                self._last_stack_action_max_prob = detached_stack[:, 4].mean()
+                self._last_stack_expected_depth = detached_stack[:, 5].mean()
+                self._last_stack_mask_fill_fraction = detached_stack[:, 6].mean()
+                self._last_stack_top_occupancy = detached_stack[:, 7].mean()
+                self._last_stack_gate_entropy = gate_entropy
+                self._last_stack_gate_entropy_normalized = gate_entropy / math.log(
+                    float(self.config.stack_slots)
+                )
+                self._last_stack_gate_max_prob = detached_stack[:, 9].mean()
+                self._last_stack_gate_effective_slots = detached_stack[:, 10].mean()
             self._last_total_loss = loss.detach() if loss is not None else None
             logits = None
         else: