cell.py

"""One physical recurrent Dendro cell containing the complete model mechanism."""

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Any

import torch
from torch.nn import functional as F

from ._source_bound import SourceBoundModule
from .cache import DendroKVCache
from .configuration_dendro_omni import DendroOmniConfig
from .modalities import DendroModalityLayout
from .source import DendroSourceLayer
from .spatial import apply_rotary_position_embedding


@dataclass(slots=True)
class DendroCellState:
    depth_index: int
    phase: str
    activation_heat: torch.Tensor
    entropy_pressure: torch.Tensor
    novelty: torch.Tensor
    salience: torch.Tensor
    route_probs: torch.Tensor
    expert_probs: torch.Tensor
    coherence: torch.Tensor
    residual_gate: torch.Tensor
    memory_write_strength: torch.Tensor
    workspace_write_strength: torch.Tensor
    plasticity_rate: torch.Tensor
    readiness: torch.Tensor
    contradiction: torch.Tensor
    attention_entropy: torch.Tensor | None = None
    top_attention_indices: torch.Tensor | None = None

    def summary(self) -> dict[str, float | int | str]:
        def mean(value: torch.Tensor) -> float:
            return float(value.detach().float().mean().cpu().item())

        return {
            "depth_index": self.depth_index,
            "phase": self.phase,
            "activation_heat": mean(self.activation_heat),
            "entropy_pressure": mean(self.entropy_pressure),
            "novelty": mean(self.novelty),
            "salience": mean(self.salience),
            "coherence": mean(self.coherence),
            "residual_gate": mean(self.residual_gate),
            "memory_write_strength": mean(self.memory_write_strength),
            "workspace_write_strength": mean(self.workspace_write_strength),
            "plasticity_rate": mean(self.plasticity_rate),
            "readiness": mean(self.readiness),
            "contradiction": mean(self.contradiction),
        }


@dataclass(slots=True)
class DendroCellOutput:
    hidden_states: torch.Tensor
    cache: DendroKVCache | None
    state: DendroCellState
    attention_weights: torch.Tensor | None = None


class DendroRecurrentCell(SourceBoundModule):
    """The single physical layer recurrently applied at all virtual depths.

    Attention, local structure, climate control, sticky plasticity, associative
    retrieval, global workspace, memory organs, routed shared-FFN computation,
    entropy regulation, coherence and dream/reflection behavior are all functions of
    one source layer.  This class contains no ``Parameter``, ``Linear`` or
    ``Embedding`` of its own.
    """

    def __init__(self, config: DendroOmniConfig, source: DendroSourceLayer) -> None:
        super().__init__(source)
        self.config = config

    def _split_heads(self, tensor: torch.Tensor) -> torch.Tensor:
        batch, length, _hidden = tensor.shape
        return tensor.view(batch, length, self.config.num_attention_heads, self.config.head_dim).transpose(1, 2)

    def _merge_heads(self, tensor: torch.Tensor) -> torch.Tensor:
        return tensor.transpose(1, 2).contiguous().flatten(-2)

    def _effort_condition(
        self,
        hidden: torch.Tensor,
        depth_code: torch.Tensor,
        phase_code: torch.Tensor,
        *,
        effort_id: int,
        effort_level: float,
        phase_progress: float,
        remaining_budget_fraction: float,
    ) -> torch.Tensor | None:
        """Build a source-derived effort/budget code without private parameters.

        The zero-strength branch intentionally requests no new source primitives,
        preserving both legacy checkpoint numerics and inference cost.
        """

        strength = float(self.config.reasoning_effort_conditioning_strength)
        if strength <= 0.0:
            return None
        batch = hidden.shape[0]
        effort_count = int(self.config.reasoning_effort_condition_count)
        checked_effort_id = min(max(0, int(effort_id)), effort_count - 1)
        # Reuse otherwise-idle rows at the end of the existing recurrent-depth
        # table.  Effort is therefore a phenotype of the same depth substrate,
        # not a separately materialized logical embedding.
        effort_row_start = max(0, int(self.config.max_recurrent_depth) - effort_count)
        effort_ids = torch.full(
            (batch,),
            effort_row_start + checked_effort_id,
            device=hidden.device,
            dtype=torch.long,
        )
        categorical = self.source.embedding(
            effort_ids,
            "recurrence/depth",
            self.config.max_recurrent_depth,
            self.config.hidden_size,
        ).unsqueeze(1)
        level = min(1.0, max(0.0, float(effort_level)))
        progress = min(1.0, max(0.0, float(phase_progress)))
        remaining = min(1.0, max(0.0, float(remaining_budget_fraction)))
        # Continuous budget information modulates already-computed depth and
        # phase codes.  This preserves their influence and source gradients while
        # avoiding another HxH projection and its saved autograd state.
        return torch.tanh(
            categorical * (0.50 + 0.50 * level + 0.25 * remaining)
            + phase_code * (0.25 + 0.50 * progress)
            + depth_code * (0.25 * remaining)
        )

    def _depth_condition(
        self,
        hidden: torch.Tensor,
        depth_idx: int,
        phase: str,
        *,
        effort_id: int,
        effort_level: float,
        phase_progress: float,
        remaining_budget_fraction: float,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        source = self.source
        batch = hidden.shape[0]
        depth_ids = torch.full((batch,), depth_idx, device=hidden.device, dtype=torch.long)
        depth_code = source.embedding(
            depth_ids,
            "recurrence/depth",
            self.config.max_recurrent_depth,
            self.config.hidden_size,
        ).unsqueeze(1)
        phase_id = {"base": 0, "reasoning": 1, "verification": 2, "dream": 3}.get(phase, 0)
        phase_ids = torch.full((batch,), phase_id, device=hidden.device, dtype=torch.long)
        phase_code = source.embedding(phase_ids, "recurrence/phase", 4, self.config.hidden_size).unsqueeze(1)
        effort_condition = self._effort_condition(
            hidden,
            depth_code,
            phase_code,
            effort_id=effort_id,
            effort_level=effort_level,
            phase_progress=phase_progress,
            remaining_budget_fraction=remaining_budget_fraction,
        )
        conditioning_code = depth_code + phase_code
        if effort_condition is not None:
            effort_strength = float(self.config.reasoning_effort_conditioning_strength)
            conditioning_code = conditioning_code + effort_strength * effort_condition
        # Depth, phase and effort all share the same FiLM transform.  The
        # zero-strength branch receives the exact legacy input and operation order.
        scale, shift = source.project_many(
            conditioning_code,
            (
                ("recurrence/film_scale", self.config.hidden_size, False),
                ("recurrence/film_shift", self.config.hidden_size, False),
            ),
        )
        conditioned = (
            hidden * (1.0 + 0.10 * torch.tanh(scale))
            + 0.10 * shift
            + 0.10 * depth_code
            + 0.05 * phase_code
        )
        if effort_condition is not None:
            conditioned = (
                conditioned
                + 0.05 * effort_strength * effort_condition
            )
        return conditioned, effort_condition

    def _context_mean(self, hidden: torch.Tensor, layout: DendroModalityLayout) -> torch.Tensor:
        """Return a mask-correct context mean without future-text leakage.

        Prefix tokens may use the complete perceptual prefix under ``prefix_bidi``;
        causal text tokens only use valid physical positions up to themselves.
        """

        valid = layout.attention_mask.unsqueeze(-1).to(hidden.dtype)
        cumulative = (hidden * valid).cumsum(dim=1)
        cumulative_count = valid.cumsum(dim=1).clamp_min(1.0)
        causal_mean = cumulative / cumulative_count
        if self.config.attention_mode == "causal":
            return causal_mean

        if self.config.attention_mode == "bidirectional":
            global_mean = (hidden * valid).sum(dim=1, keepdim=True) / valid.sum(dim=1, keepdim=True).clamp_min(1.0)
            return global_mean.expand_as(hidden)

        prefix_valid = (layout.is_prefix & layout.attention_mask).unsqueeze(-1)
        prefix_weight = prefix_valid.to(hidden.dtype)
        prefix_mean = (hidden * prefix_weight).sum(dim=1, keepdim=True)
        prefix_mean = prefix_mean / prefix_weight.sum(dim=1, keepdim=True).clamp_min(1.0)
        return torch.where(prefix_valid, prefix_mean.expand_as(hidden), causal_mean)

    def _climate(
        self,
        hidden: torch.Tensor,
        *,
        layout: DendroModalityLayout,
        depth_idx: int,
        cache: DendroKVCache | None,
    ) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
        source = self.source
        activation_heat = hidden.float().pow(2).mean(dim=-1, keepdim=True).to(hidden.dtype)
        feature_probs = torch.softmax(hidden.float(), dim=-1)
        entropy = -(feature_probs * feature_probs.clamp_min(1e-9).log()).sum(dim=-1, keepdim=True)
        entropy = (entropy / math.log(max(2, hidden.shape[-1]))).to(hidden.dtype)
        centered = hidden - self._context_mean(hidden, layout)
        novelty = centered.float().pow(2).mean(dim=-1, keepdim=True).clamp_min(1e-12).sqrt().to(hidden.dtype)
        memory_pressure = torch.zeros_like(activation_heat)
        route_imbalance = torch.zeros_like(activation_heat)
        plasticity_volatility = torch.zeros_like(activation_heat)
        # Runtime organs are consumed through causal scans below. Feeding their
        # *final* cached summaries back into every token here would make chunked
        # decoding differ from full-sequence training. These slots remain reserved
        # for source-compatible climate extensions that provide tokenwise histories.
        del cache, depth_idx
        metrics = torch.cat(
            [activation_heat, entropy, novelty, memory_pressure, route_imbalance, plasticity_volatility],
            dim=-1,
        )
        controls_raw = source.project(metrics, "climate/controller", 8, low_bit=False)
        controls = {
            "temperature": 0.55 + 0.90 * torch.sigmoid(controls_raw[..., 0:1]),
            "residual_gate": 0.10 + 0.90 * torch.sigmoid(controls_raw[..., 1:2]),
            "plasticity_rate": 0.20 * torch.sigmoid(controls_raw[..., 2:3]),
            "memory_write": torch.sigmoid(controls_raw[..., 3:4]),
            "workspace_write": torch.sigmoid(controls_raw[..., 4:5]),
            "attention_focus": torch.sigmoid(controls_raw[..., 5:6]),
            "entropy_compress": torch.sigmoid(controls_raw[..., 6:7]),
            "dream_gate": torch.sigmoid(controls_raw[..., 7:8]),
        }
        return {
            "activation_heat": activation_heat,
            "entropy": entropy,
            "novelty": novelty,
            "memory_pressure": memory_pressure,
            "route_imbalance": route_imbalance,
            "plasticity_volatility": plasticity_volatility,
        }, controls

    def _plasticity(
        self,
        hidden: torch.Tensor,
        controls: dict[str, torch.Tensor],
        layout: DendroModalityLayout,
        depth_idx: int,
        cache: DendroKVCache | None,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        source = self.source
        salience = source.gate(hidden, "plasticity/salience", 1)
        if hidden.shape[1] > 1:
            previous = F.pad(hidden[:, :-1], (0, 0, 1, 0))
            novelty = 1.0 - F.cosine_similarity(hidden.float(), previous.float(), dim=-1).unsqueeze(-1)
            novelty[:, 0] = 0.0
            novelty = novelty.to(hidden.dtype)
        else:
            novelty = torch.zeros_like(salience)
        old_trace = cache.get_runtime_state("plasticity_trace", depth_idx) if cache is not None else None
        if old_trace is None:
            old_trace = torch.zeros(hidden.shape[0], hidden.shape[-1], device=hidden.device, dtype=hidden.dtype)
        else:
            old_trace = old_trace.to(device=hidden.device, dtype=hidden.dtype)
        writes = salience * (1.0 + novelty) * hidden
        rates = controls["plasticity_rate"]
        valid = layout.attention_mask.unsqueeze(-1)
        trace = old_trace
        token_traces: list[torch.Tensor] = []
        # This is an actual sticky causal state scan.  It makes the full-sequence
        # training path obey the same no-future contract as token-by-token decoding.
        for token_idx in range(hidden.shape[1]):
            candidate = self.config.plasticity_decay * trace + rates[:, token_idx] * writes[:, token_idx]
            trace = torch.where(valid[:, token_idx], candidate, trace)
            token_traces.append(trace)
        stacked = torch.stack(token_traces, dim=1)
        modulation = source.project(stacked, "plasticity/trace_modulation", self.config.hidden_size, low_bit=False)
        return salience, novelty, stacked + 0.05 * modulation, trace

    def _make_attention_mask(
        self,
        *,
        query_layout: DendroModalityLayout,
        key_positions: torch.Tensor,
        key_is_prefix: torch.Tensor,
        key_attention_mask: torch.Tensor,
    ) -> torch.Tensor:
        q_pos = query_layout.sequence_positions.unsqueeze(-1)
        k_pos = key_positions.unsqueeze(-2)
        q_prefix = query_layout.is_prefix.unsqueeze(-1)
        k_prefix = key_is_prefix.unsqueeze(-2)
        mode = self.config.attention_mode
        if mode == "bidirectional":
            allowed = torch.ones_like(q_pos <= k_pos, dtype=torch.bool)
        elif mode == "causal":
            allowed = k_pos <= q_pos
        else:  # prefix_bidi
            allowed = (q_prefix & k_prefix) | (~q_prefix & (k_prefix | (k_pos <= q_pos)))
        q_valid = query_layout.attention_mask.unsqueeze(-1)
        k_valid = key_attention_mask.unsqueeze(-2)
        allowed = allowed & q_valid & k_valid
        # SDPA rows may not be entirely masked. Invalid query rows are later zeroed.
        first_key = torch.zeros_like(allowed)
        first_key[..., 0] = True
        allowed = allowed | (~q_valid & first_key)
        return allowed

    def _make_local_attention_mask(
        self,
        *,
        query_layout: DendroModalityLayout,
        key_positions: torch.Tensor,
        key_is_prefix: torch.Tensor,
        key_attention_mask: torch.Tensor,
    ) -> torch.Tensor:
        """Cache-aware three-position local mask over the shared Q/K/V stream."""

        full = self._make_attention_mask(
            query_layout=query_layout,
            key_positions=key_positions,
            key_is_prefix=key_is_prefix,
            key_attention_mask=key_attention_mask,
        )
        q_pos = query_layout.sequence_positions.unsqueeze(-1)
        k_pos = key_positions.unsqueeze(-2)
        distance = q_pos - k_pos
        if self.config.attention_mode == "bidirectional":
            local = distance.abs() <= 1
        elif self.config.attention_mode == "causal":
            local = (distance >= 0) & (distance < 3)
        else:
            q_prefix = query_layout.is_prefix.unsqueeze(-1)
            k_prefix = key_is_prefix.unsqueeze(-2)
            prefix_local = k_prefix & (distance.abs() <= 1)
            text_local = (distance >= 0) & (distance < 3)
            local = torch.where(q_prefix, prefix_local, text_local)
        local = local & full
        q_valid = query_layout.attention_mask.unsqueeze(-1)
        first_key = torch.zeros_like(local)
        first_key[..., 0] = True
        return local | (~q_valid & first_key)

    def _attention(
        self,
        hidden: torch.Tensor,
        *,
        layout: DendroModalityLayout,
        controls: dict[str, torch.Tensor],
        plasticity_trace: torch.Tensor,
        depth_idx: int,
        cache: DendroKVCache | None,
        use_cache: bool,
        output_attentions: bool,
    ) -> tuple[
        torch.Tensor,
        torch.Tensor,
        torch.Tensor | None,
        torch.Tensor,
        tuple[torch.Tensor | None, torch.Tensor | None],
    ]:
        source = self.source
        qkv = source.project(hidden, "cell/attention/qkv", 3 * self.config.hidden_size)
        q, k, v = qkv.chunk(3, dim=-1)
        q, k, v = self._split_heads(q), self._split_heads(k), self._split_heads(v)
        if self.config.qkv_norm:
            if self.config.align_qkv_norms:
                q, k, v = source.aligned_qkv_norm(q, k, v, "cell/attention/qkv_norm", eps=self.config.layer_norm_eps)
            else:
                q = source.rms_norm(q, "cell/attention/q_norm", eps=self.config.layer_norm_eps)
                k = source.rms_norm(k, "cell/attention/k_norm", eps=self.config.layer_norm_eps)
                v = source.rms_norm(v, "cell/attention/v_norm", eps=self.config.layer_norm_eps)
        q, k = apply_rotary_position_embedding(
            q,
            k,
            layout.sequence_positions,
            theta=self.config.rope_theta,
        )
        trace_heads = plasticity_trace.view(
            plasticity_trace.shape[0],
            plasticity_trace.shape[1],
            self.config.num_attention_heads,
            self.config.head_dim,
        ).transpose(1, 2)
        q = q + 0.02 * trace_heads

        if use_cache:
            if cache is None:
                raise RuntimeError("use_cache=True requires a DendroKVCache")
            key, value = cache.update(
                k,
                v,
                depth_idx,
                {
                    "is_prefix": layout.is_prefix,
                    "positions": layout.sequence_positions,
                    "attention_mask": layout.attention_mask,
                    "modality_ids": layout.modality_ids,
                },
            )
            key_positions = cache.key_positions
            key_is_prefix = cache.key_is_prefix
            key_attention = cache.key_attention_mask
            assert key_positions is not None and key_is_prefix is not None and key_attention is not None
            key_positions = key_positions.to(hidden.device)
            key_is_prefix = key_is_prefix.to(hidden.device)
            key_attention = key_attention.to(hidden.device)
        else:
            key, value = k, v
            key_positions = layout.sequence_positions
            key_is_prefix = layout.is_prefix
            key_attention = layout.attention_mask

        allowed = self._make_attention_mask(
            query_layout=layout,
            key_positions=key_positions,
            key_is_prefix=key_is_prefix,
            key_attention_mask=key_attention,
        )
        attn_mask = allowed.unsqueeze(1)
        local_mask = self._make_local_attention_mask(
            query_layout=layout,
            key_positions=key_positions,
            key_is_prefix=key_is_prefix,
            key_attention_mask=key_attention,
        ).unsqueeze(1)
        dropout_p = self.config.attention_dropout if self.training else 0.0
        attention_weights = None
        attention_entropy = None
        top_indices = None
        scale = 1.0 / math.sqrt(self.config.head_dim)
        temperature = controls["temperature"].transpose(1, 2).unsqueeze(-1).to(q.dtype)
        tempered_q = q / temperature

        if output_attentions:
            logits = torch.matmul(tempered_q.float(), key.float().transpose(-1, -2)) * scale
            logits = logits.masked_fill(~attn_mask, torch.finfo(logits.dtype).min)
            attention_weights = torch.softmax(logits, dim=-1).to(hidden.dtype)
            attention_weights = F.dropout(attention_weights, p=dropout_p, training=self.training)
            context = torch.matmul(attention_weights, value)
            probs = attention_weights.float().clamp_min(1e-9)
            attention_entropy = -(probs * probs.log()).sum(dim=-1).mean(dim=1)
            top_indices = attention_weights.detach().mean(dim=1).topk(
                k=min(4, attention_weights.shape[-1]), dim=-1
            ).indices
        else:
            context = F.scaled_dot_product_attention(
                tempered_q,
                key,
                value,
                attn_mask=attn_mask,
                dropout_p=dropout_p,
                is_causal=False,
                scale=scale,
            )
        local_context = F.scaled_dot_product_attention(
            tempered_q,
            key,
            value,
            attn_mask=local_mask,
            dropout_p=dropout_p,
            is_causal=False,
            scale=scale,
        )
        context = context * layout.attention_mask[:, None, :, None].to(context.dtype)
        local_context = local_context * layout.attention_mask[:, None, :, None].to(local_context.dtype)

        # Shared head-communication state lets heads exchange summaries without a
        # second attention module or independent parameters.
        head_summary = self._merge_heads(context)
        previous_comm = cache.get_runtime_state("head_communication", depth_idx) if cache is not None else None
        if previous_comm is None:
            comm_state = torch.zeros(
                head_summary.shape[0],
                self.config.hidden_size,
                device=head_summary.device,
                dtype=head_summary.dtype,
            )
        else:
            comm_state = previous_comm.to(head_summary.device, head_summary.dtype)
        comm_tokens: list[torch.Tensor] = []
        for token_idx in range(head_summary.shape[1]):
            candidate = source.project(
                head_summary[:, token_idx] + comm_state,
                "cell/attention/head_communication",
                self.config.hidden_size,
                low_bit=False,
            )
            valid = layout.attention_mask[:, token_idx, None]
            comm_state = torch.where(valid, candidate, comm_state)
            comm_tokens.append(torch.where(valid, candidate, torch.zeros_like(candidate)))
        comm = torch.stack(comm_tokens, dim=1)
        comm_gate = source.gate(hidden, "cell/attention/head_communication_gate", self.config.hidden_size)
        comm_heads = self._split_heads(comm)
        gate_heads = self._split_heads(comm_gate)
        context = context + 0.05 * comm_heads * gate_heads
        return (
            self._merge_heads(context),
            self._merge_heads(local_context),
            attention_weights,
            comm_state,
            (attention_entropy, top_indices),
        )

    def _memory_slots(
        self,
        hidden: torch.Tensor,
        cache: DendroKVCache | None,
        depth_idx: int,
    ) -> torch.Tensor:
        cached_memory = cache.get_runtime_state("memory", depth_idx) if cache is not None else None
        if cached_memory is not None:
            return cached_memory.to(device=hidden.device, dtype=hidden.dtype)
        seeds = self.source.primitive("memory/slots", (self.config.memory_slots, self.config.hidden_size))
        organ_ids = torch.arange(self.config.memory_slots, device=hidden.device) % self.config.num_memory_organs
        organs = self.source.embedding(
            organ_ids,
            "memory/organs",
            self.config.num_memory_organs,
            self.config.hidden_size,
        )
        return (seeds + 0.10 * organs).unsqueeze(0).expand(hidden.shape[0], -1, -1)

    def _workspace_slots(
        self,
        hidden: torch.Tensor,
        cache: DendroKVCache | None,
        depth_idx: int,
    ) -> torch.Tensor:
        cached_workspace = cache.get_runtime_state("workspace", depth_idx) if cache is not None else None
        if cached_workspace is not None:
            return cached_workspace.to(device=hidden.device, dtype=hidden.dtype)
        seeds = self.source.primitive("workspace/slots", (self.config.workspace_slots, self.config.hidden_size))
        return seeds.unsqueeze(0).expand(hidden.shape[0], -1, -1)

    @staticmethod
    def _affine_slot_states(
        initial: torch.Tensor,
        updates: torch.Tensor,
        valid: torch.Tensor,
        decay: float,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        """Return every pre-update state and the final affine recurrent state."""

        scan_dtype = torch.float32 if initial.dtype in {torch.float16, torch.bfloat16} else initial.dtype
        scan_valid = valid.unsqueeze(-1).unsqueeze(-1)
        multiplier = torch.where(
            scan_valid,
            torch.full_like(scan_valid, float(decay), dtype=scan_dtype),
            torch.ones_like(scan_valid, dtype=scan_dtype),
        )
        additive = updates.to(scan_dtype) * scan_valid
        products = torch.cumprod(multiplier, dim=1)
        scaled = additive / products.clamp_min(torch.finfo(scan_dtype).tiny)
        inclusive = torch.cumsum(scaled, dim=1)
        before_sum = inclusive - scaled
        before_product = torch.cat([torch.ones_like(products[:, :1]), products[:, :-1]], dim=1)
        initial_scan = initial.to(scan_dtype)
        states_before = before_product * (initial_scan.unsqueeze(1) + before_sum)
        final = products[:, -1] * (initial_scan + inclusive[:, -1])
        return states_before.to(initial.dtype), final.to(initial.dtype)

    def _slot_scan(
        self,
        hidden: torch.Tensor,
        slots: torch.Tensor,
        controls: dict[str, torch.Tensor],
        layout: DendroModalityLayout,
        name: str,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """Read then write a memory/workspace organ token by token.

        This is a real causal state machine: token ``t`` reads only the organ state
        created by cached history and tokens ``< t``, then writes the state used by
        token ``t+1``.  The same routine is used in full-sequence training and
        one-token cached generation.
        """

        if name not in {"memory", "workspace"}:
            raise ValueError(f"unsupported slot scan {name!r}")
        source = self.source
        decay = self.config.memory_decay if name == "memory" else self.config.workspace_decay
        strength_key = "memory_write" if name == "memory" else "workspace_write"
        # Token projections are independent of the recurrent slot state. Project
        # the complete sequence once; only the actual read/write state transition
        # remains causal below. The affine recurrence itself has a closed-form
        # prefix scan, evaluated in bounded blocks to avoid long-context underflow.
        queries = source.project(hidden, f"{name}/query", self.config.hidden_size, low_bit=False)
        routers = torch.softmax(
            source.project(hidden, f"{name}/write_router", slots.shape[1], low_bit=False).float(),
            dim=-1,
        ).to(hidden.dtype)
        writes = source.project(hidden, f"{name}/write_value", self.config.hidden_size, low_bit=False)
        strength = controls[strength_key]
        valid = layout.attention_mask
        if hidden.shape[1] == 1:
            key = source.project(slots, f"{name}/key", self.config.hidden_size, low_bit=False)
            value = source.project(slots, f"{name}/value", self.config.hidden_size, low_bit=False)
            scores = torch.einsum("bh,bsh->bs", queries[:, 0].float(), key.float())
            scores = scores / math.sqrt(self.config.hidden_size)
            probabilities = torch.softmax(scores, dim=-1).to(hidden.dtype)
            read = torch.einsum("bs,bsh->bh", probabilities, value).unsqueeze(1)
            update = strength[:, 0].unsqueeze(1) * routers[:, 0].unsqueeze(-1) * writes[:, 0].unsqueeze(1)
            candidate = decay * slots + update
            mask = valid[:, 0, None, None]
            final = torch.where(mask, candidate, slots)
            final_router = torch.where(valid[:, 0, None], routers[:, 0], torch.zeros_like(routers[:, 0]))
            return read, final, final_router
        updates = (
            strength.unsqueeze(-1)
            * routers.unsqueeze(-1)
            * writes.unsqueeze(-2)
        )

        # A 512-token training window is deliberately handled by one tensor scan,
        # without entering a Python block loop. Extremely long contexts retain a
        # bounded fallback so products cannot underflow and peak memory stays sane.
        scan_block = 1024
        if hidden.shape[1] <= scan_block:
            states_before, current = self._affine_slot_states(slots, updates, valid, decay)
            key = source.project(states_before, f"{name}/key", self.config.hidden_size, low_bit=False)
            value = source.project(states_before, f"{name}/value", self.config.hidden_size, low_bit=False)
            scores = torch.einsum("bth,btsh->bts", queries.float(), key.float())
            scores = scores / math.sqrt(self.config.hidden_size)
            read_probs = torch.softmax(scores, dim=-1).to(hidden.dtype)
            reads = torch.einsum("bts,btsh->bth", read_probs, value)
        else:
            read_blocks: list[torch.Tensor] = []
            current = slots
            for start in range(0, hidden.shape[1], scan_block):
                end = min(hidden.shape[1], start + scan_block)
                states_before, current = self._affine_slot_states(
                    current,
                    updates[:, start:end],
                    valid[:, start:end],
                    decay,
                )
                key = source.project(states_before, f"{name}/key", self.config.hidden_size, low_bit=False)
                value = source.project(states_before, f"{name}/value", self.config.hidden_size, low_bit=False)
                scores = torch.einsum("bth,btsh->bts", queries[:, start:end].float(), key.float())
                scores = scores / math.sqrt(self.config.hidden_size)
                read_probs = torch.softmax(scores, dim=-1).to(hidden.dtype)
                read_blocks.append(torch.einsum("bts,btsh->bth", read_probs, value))
            reads = torch.cat(read_blocks, dim=1)

        positions = torch.arange(hidden.shape[1], device=hidden.device).unsqueeze(0)
        last_index = torch.where(valid, positions, -1).amax(dim=1)
        safe_index = last_index.clamp_min(0)
        last_router = routers.gather(
            1,
            safe_index[:, None, None].expand(-1, 1, routers.shape[-1]),
        ).squeeze(1)
        last_router = torch.where((last_index >= 0).unsqueeze(-1), last_router, torch.zeros_like(last_router))
        return reads, current, last_router

    def _associative_scan(
        self,
        hidden: torch.Tensor,
        salience: torch.Tensor,
        layout: DendroModalityLayout,
        cache: DendroKVCache | None,
        depth_idx: int,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        """Bounded causal associative recall and online write path."""

        source = self.source
        batch, _length, hidden_size = hidden.shape
        keys = cache.get_runtime_state("associative_keys", depth_idx) if cache is not None else None
        values = cache.get_runtime_state("associative_values", depth_idx) if cache is not None else None
        scores = cache.get_runtime_state("associative_scores", depth_idx) if cache is not None else None
        if keys is None or values is None:
            keys = hidden.new_empty(batch, 0, hidden_size)
            values = hidden.new_empty(batch, 0, hidden_size)
            scores = hidden.new_empty(batch, 0)
        else:
            keys = keys.to(device=hidden.device, dtype=hidden.dtype)
            values = values.to(device=hidden.device, dtype=hidden.dtype)
            if scores is None:
                scores = torch.ones(batch, keys.shape[1], device=hidden.device, dtype=hidden.dtype)
            else:
                scores = scores.to(device=hidden.device, dtype=hidden.dtype)

        projected_queries = source.project(hidden, "associative/query", hidden_size, low_bit=False)
        projected_keys = source.project(hidden, "associative/key", hidden_size, low_bit=False)
        projected_values = source.project(hidden, "associative/value", hidden_size, low_bit=False)
        new_scores = torch.where(
            layout.attention_mask,
            salience[..., 0],
            torch.full_like(salience[..., 0], -1.0),
        )
        if hidden.shape[1] == 1:
            if keys.shape[1] == 0:
                read = torch.zeros_like(hidden)
            else:
                logits = torch.einsum("bh,bkh->bk", projected_queries[:, 0].float(), keys.float())
                logits = logits / math.sqrt(hidden_size)
                logits = logits + scores.float().clamp_min(1e-8).log()
                probabilities = torch.softmax(logits, dim=-1).to(hidden.dtype)
                read = torch.einsum("bk,bkh->bh", probabilities, values).unsqueeze(1)
            keys = torch.cat([keys, projected_keys], dim=1)
            values = torch.cat([values, projected_values], dim=1)
            scores = torch.cat([scores, new_scores], dim=1)
            keep = min(self.config.associative_slots, keys.shape[1])
            final_scores, final_indices = scores.topk(keep, dim=1)
            gather = final_indices.unsqueeze(-1).expand(-1, -1, hidden_size)
            return read, keys.gather(1, gather), values.gather(1, gather), final_scores
        initial_count = keys.shape[1]
        candidate_keys = torch.cat([keys, projected_keys], dim=1)
        candidate_values = torch.cat([values, projected_values], dim=1)
        candidate_scores = torch.cat([scores, new_scores], dim=1)
        candidate_count = candidate_keys.shape[1]
        keep = min(self.config.associative_slots, candidate_count)

        token_index = torch.arange(hidden.shape[1], device=hidden.device).view(1, -1, 1)
        candidate_index = torch.arange(candidate_count, device=hidden.device).view(1, 1, -1)
        allowed = candidate_index < (initial_count + token_index)
        ranked = candidate_scores.unsqueeze(1).expand(-1, hidden.shape[1], -1).masked_fill(~allowed, float("-inf"))
        _top_scores, top_indices = ranked.topk(keep, dim=-1)
        # Do not expand candidates to [B, T, C, H] before gather. Although that
        # expansion is a cheap forward view, GatherBackward allocates its full
        # gradient (42+ GiB for T=3340/H=1024). Flattened batch offsets let
        # IndexSelectBackward accumulate directly into the compact [B, C, H]
        # candidate table while returning the identical [B, T, K, H] values.
        batch_offsets = (
            torch.arange(batch, device=hidden.device, dtype=top_indices.dtype)
            * candidate_count
        ).view(batch, 1, 1)
        flat_indices = (top_indices + batch_offsets).reshape(-1)
        selected_shape = (*top_indices.shape, hidden_size)
        selected_keys = candidate_keys.reshape(
            batch * candidate_count, hidden_size
        ).index_select(0, flat_indices).reshape(selected_shape)
        selected_values = candidate_values.reshape(
            batch * candidate_count, hidden_size
        ).index_select(0, flat_indices).reshape(selected_shape)
        selected_scores = candidate_scores.unsqueeze(1).expand(-1, hidden.shape[1], -1).gather(2, top_indices)
        selected_valid = allowed.expand(hidden.shape[0], -1, -1).gather(2, top_indices)

        logits = torch.einsum("bth,btkh->btk", projected_queries.float(), selected_keys.float())
        logits = logits / math.sqrt(hidden_size)
        logits = logits + selected_scores.float().clamp_min(1e-8).log()
        logits = logits.masked_fill(~selected_valid, -1e9)
        probs = torch.softmax(logits, dim=-1).to(hidden.dtype) * selected_valid.to(hidden.dtype)
        probs = probs / probs.sum(dim=-1, keepdim=True).clamp_min(1e-8)
        reads = torch.einsum("btk,btkh->bth", probs, selected_values)

        final_keep = min(self.config.associative_slots, candidate_count)
        final_scores, final_indices = candidate_scores.topk(final_keep, dim=1)
        final_gather = final_indices.unsqueeze(-1).expand(-1, -1, hidden_size)
        final_keys = candidate_keys.gather(1, final_gather)
        final_values = candidate_values.gather(1, final_gather)
        return reads, final_keys, final_values, final_scores

    def _route_mix(
        self,
        hidden: torch.Tensor,
        attention: torch.Tensor,
        local_attention: torch.Tensor,
        memory: torch.Tensor,
        workspace: torch.Tensor,
        associative: torch.Tensor,
        controls: dict[str, torch.Tensor],
        layout: DendroModalityLayout,
        effort_condition: torch.Tensor | None = None,
        *,
        effort_level: float = 0.0,
        phase: str = "base",
        phase_progress: float = 1.0,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        source = self.source
        del layout
        local = source.project(local_attention, "routes/local", self.config.hidden_size, low_bit=False)
        residual_route = source.project(hidden, "routes/residual", self.config.hidden_size, low_bit=False)
        components = [attention, local, memory, workspace, associative, residual_route]
        while len(components) < self.config.num_routes:
            index = len(components)
            components.append(source.project(hidden, f"routes/aux_{index}", self.config.hidden_size, low_bit=False))
        components = components[: self.config.num_routes]
        route_logits = source.project(hidden, "routes/router", self.config.num_routes, low_bit=False)
        if effort_condition is not None:
            route_logits = route_logits + float(
                self.config.reasoning_effort_conditioning_strength
            ) * source.project(
                effort_condition,
                "routes/router",
                self.config.num_routes,
                bias=False,
                low_bit=False,
            )
        prior_strength = float(self.config.reasoning_route_prior_strength)
        if prior_strength > 0.0 and self.config.num_routes > 0:
            # Route order is attention, local, memory, workspace, associative,
            # residual.  Deliberation should increasingly consult shared memory
            # and workspace rather than repeatedly amplifying the residual path.
            # This is an architectural prior, not a task-answer heuristic, and it
            # introduces no parameters or checkpoint memory.
            semantic_prior = hidden.new_tensor(
                [0.25, -0.25, 0.35, 0.55, 0.45, -0.55]
            )
            if self.config.num_routes < semantic_prior.numel():
                semantic_prior = semantic_prior[: self.config.num_routes]
            elif self.config.num_routes > semantic_prior.numel():
                semantic_prior = F.pad(
                    semantic_prior,
                    (0, self.config.num_routes - semantic_prior.numel()),
                )
            semantic_prior = semantic_prior - semantic_prior.mean()
            level = min(1.0, max(0.0, float(effort_level)))
            progress = min(1.0, max(0.0, float(phase_progress)))
            if phase == "base":
                phase_gain = 0.25 * level
            elif phase == "reasoning":
                phase_gain = (0.75 + 0.25 * level) * (0.75 + 0.25 * progress)
            elif phase == "verification":
                phase_gain = 1.0
            else:
                phase_gain = 0.50 * level
            route_logits = route_logits + prior_strength * phase_gain * semantic_prior
        focus = controls["attention_focus"]
        if self.config.num_routes > 0:
            route_logits[..., :1] = route_logits[..., :1] + focus
        route_probs = torch.softmax(route_logits.float(), dim=-1).to(hidden.dtype)
        stacked = torch.stack(components, dim=-2)
        mixed = (route_probs.unsqueeze(-1) * stacked).sum(dim=-2)
        return mixed, route_probs

    def _shared_routed_ffn(
        self,
        hidden: torch.Tensor,
        effort_condition: torch.Tensor | None = None,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        source = self.source
        normalized = source.rms_norm(hidden, "cell/ffn/input_norm", eps=self.config.layer_norm_eps)
        gate, value = source.project_many(
            normalized,
            (
                ("cell/ffn/gate", self.config.intermediate_size, None),
                ("cell/ffn/value", self.config.intermediate_size, None),
            ),
        )
        router_logits = source.project(normalized, "cell/ffn/expert_router", self.config.num_experts, low_bit=False)
        if effort_condition is not None:
            router_logits = router_logits + float(
                self.config.reasoning_effort_conditioning_strength
            ) * source.project(
                effort_condition,
                "cell/ffn/expert_router",
                self.config.num_experts,
                bias=False,
                low_bit=False,
            )
        probs = torch.softmax(router_logits.float(), dim=-1).to(hidden.dtype)
        if self.config.expert_top_k < self.config.num_experts:
            top_values, top_indices = probs.topk(self.config.expert_top_k, dim=-1)
            sparse = torch.zeros_like(probs).scatter(-1, top_indices, top_values)
            probs = sparse / sparse.sum(dim=-1, keepdim=True).clamp_min(1e-8)
        # Experts are source-derived channel phenotypes over one shared FFN, not
        # duplicated expert matrices.
        expert_codes = self.source.primitive(
            "cell/ffn/expert_codes",
            (self.config.num_experts, self.config.intermediate_size),
        )
        modulation = torch.matmul(probs, expert_codes)
        activated = F.silu(gate) * value * (1.0 + 0.15 * torch.tanh(modulation))
        output = source.project(activated, "cell/ffn/down", self.config.hidden_size)
        return output, probs

    def _coherence_and_entropy(
        self,
        hidden: torch.Tensor,
        proposal: torch.Tensor,
        controls: dict[str, torch.Tensor],
        *,
        phase: str,
    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
        source = self.source
        identity = source.project(hidden, "coherence/identity", self.config.hidden_size, low_bit=False)
        whole = source.project(proposal, "coherence/whole", self.config.hidden_size, low_bit=False)
        coherence = F.cosine_similarity(identity.float(), whole.float(), dim=-1).unsqueeze(-1).to(hidden.dtype)
        residual_gate = controls["residual_gate"] * torch.sigmoid(2.0 * coherence)
        merged = hidden + self.config.residual_scale * residual_gate * proposal

        normalized = source.rms_norm(merged, "entropy/input_norm", eps=self.config.layer_norm_eps)
        compressed = source.project(normalized, "entropy/compression", self.config.hidden_size, low_bit=False)
        merged = merged + 0.10 * controls["entropy_compress"] * torch.tanh(compressed)

        # Dream/reflection is deterministic and source-derived. It introduces no
        # random inference drift and remains cache-parity friendly.
        dream = source.project(torch.sin(normalized), "dream/reflection", self.config.hidden_size, low_bit=False)
        phase_strength = 1.0 if phase in {"reasoning", "verification", "dream"} else 0.25
        merged = merged + 0.05 * phase_strength * controls["dream_gate"] * torch.tanh(dream)
        readiness = source.gate(torch.cat([merged, whole], dim=-1), "reasoning/readiness", 1)
        contradiction = source.gate(torch.cat([merged, -whole], dim=-1), "reasoning/contradiction", 1)
        correction_strength = float(self.config.reasoning_correction_strength)
        if correction_strength > 0.0:
            # A signed gate is neutral when both uncalibrated heads sit at 0.5.
            # Once trained, readiness advances a proposal while contradiction
            # suppresses or reverses it.  The feature is opt-in for compatibility.
            correction_gate = (readiness - contradiction).clamp(-1.0, 1.0)
            correction_phase = 1.0 if phase in {"reasoning", "verification"} else 0.25
            merged = (
                merged
                + correction_strength
                * correction_phase
                * correction_gate
                * torch.tanh(proposal)
            )
        return merged, coherence, readiness, contradiction

    def _update_runtime_state(
        self,
        *,
        memory: torch.Tensor,
        memory_scores: torch.Tensor,
        workspace: torch.Tensor,
        associative_keys: torch.Tensor,
        associative_values: torch.Tensor,
        associative_scores: torch.Tensor,
        route_probs: torch.Tensor,
        plasticity_trace: torch.Tensor,
        communication: torch.Tensor,
        depth_idx: int,
        cache: DendroKVCache,
    ) -> None:
        # The tokenwise scans already produced the exact causal final states. This
        # method only commits them to the cache; it performs no second hidden update.
        cache.set_runtime_state("memory", memory, depth_idx)
        cache.set_runtime_state("memory_scores", memory_scores, depth_idx)
        cache.set_runtime_state("workspace", workspace, depth_idx)
        cache.set_runtime_state("plasticity_trace", plasticity_trace, depth_idx)
        cache.set_runtime_state("head_communication", communication, depth_idx)

        route_mean = route_probs.mean(dim=1)
        old_route = cache.get_runtime_state("route_history", depth_idx)
        route_history = route_mean if old_route is None else 0.90 * old_route.to(route_mean.device) + 0.10 * route_mean
        cache.set_runtime_state("route_history", route_history, depth_idx)
        cache.set_runtime_state("associative_keys", associative_keys, depth_idx)
        cache.set_runtime_state("associative_values", associative_values, depth_idx)
        cache.set_runtime_state("associative_scores", associative_scores, depth_idx)

    def forward(
        self,
        hidden_states: torch.Tensor,
        *,
        layout: DendroModalityLayout,
        depth_idx: int,
        phase: str = "base",
        effort_id: int = 0,
        effort_level: float = 0.0,
        phase_progress: float = 1.0,
        remaining_budget_fraction: float = 0.0,
        cache: DendroKVCache | None = None,
        use_cache: bool = False,
        output_attentions: bool = False,
    ) -> DendroCellOutput:
        if hidden_states.ndim != 3 or hidden_states.shape[-1] != self.config.hidden_size:
            raise ValueError("hidden_states must be [batch, sequence, hidden_size]")
        if not 0 <= depth_idx < self.config.max_recurrent_depth:
            raise ValueError("depth_idx exceeds max_recurrent_depth")
        source = self.source
        conditioned, effort_condition = self._depth_condition(
            hidden_states,
            depth_idx,
            phase,
            effort_id=effort_id,
            effort_level=effort_level,
            phase_progress=phase_progress,
            remaining_budget_fraction=remaining_budget_fraction,
        )
        normalized = source.rms_norm(conditioned, "cell/input_norm", eps=self.config.layer_norm_eps)
        climate, controls = self._climate(
            normalized,
            layout=layout,
            depth_idx=depth_idx,
            cache=cache,
        )
        salience, novelty, token_trace, final_trace = self._plasticity(
            normalized,
            controls,
            layout,
            depth_idx,
            cache,
        )

        attention, local_attention, attention_weights, communication, diagnostics = self._attention(
            normalized,
            layout=layout,
            controls=controls,
            plasticity_trace=token_trace,
            depth_idx=depth_idx,
            cache=cache,
            use_cache=use_cache,
            output_attentions=output_attentions,
        )
        memory_slots = self._memory_slots(normalized, cache, depth_idx)
        workspace_slots = self._workspace_slots(normalized, cache, depth_idx)
        memory_read, final_memory, memory_scores = self._slot_scan(
            normalized,
            memory_slots,
            controls,
            layout,
            "memory",
        )
        workspace_read, final_workspace, _workspace_scores = self._slot_scan(
            normalized,
            workspace_slots,
            controls,
            layout,
            "workspace",
        )
        associative_read, associative_keys, associative_values, associative_scores = self._associative_scan(
            normalized,
            salience,
            layout,
            cache,
            depth_idx,
        )
        mixed, route_probs = self._route_mix(
            normalized,
            attention,
            local_attention,
            memory_read,
            workspace_read,
            associative_read,
            controls,
            layout,
            effort_condition,
            effort_level=effort_level,
            phase=phase,
            phase_progress=phase_progress,
        )
        attention_out = source.project(mixed, "cell/attention/output", self.config.hidden_size)
        hidden = conditioned + self.config.residual_scale * controls["residual_gate"] * attention_out
        ffn_out, expert_probs = self._shared_routed_ffn(hidden, effort_condition)
        hidden, coherence, readiness, contradiction = self._coherence_and_entropy(
            hidden,
            ffn_out,
            controls,
            phase=phase,
        )
        hidden = F.dropout(hidden, p=self.config.dropout, training=self.training)
        hidden = hidden * layout.attention_mask.unsqueeze(-1).to(hidden.dtype)

        if use_cache:
            assert cache is not None
            self._update_runtime_state(
                memory=final_memory,
                memory_scores=memory_scores,
                workspace=final_workspace,
                associative_keys=associative_keys,
                associative_values=associative_values,
                associative_scores=associative_scores,
                route_probs=route_probs,
                plasticity_trace=final_trace,
                communication=communication,
                depth_idx=depth_idx,
                cache=cache,
            )

        attention_entropy, top_indices = diagnostics if diagnostics is not None else (None, None)
        state = DendroCellState(
            depth_index=depth_idx,
            phase=phase,
            activation_heat=climate["activation_heat"],
            entropy_pressure=climate["entropy"],
            novelty=novelty,
            salience=salience,
            route_probs=route_probs,
            expert_probs=expert_probs,
            coherence=coherence,
            residual_gate=controls["residual_gate"],
            memory_write_strength=controls["memory_write"],
            workspace_write_strength=controls["workspace_write"],
            plasticity_rate=controls["plasticity_rate"],
            readiness=readiness,
            contradiction=contradiction,
            attention_entropy=attention_entropy,
            top_attention_indices=top_indices,
        )
        return DendroCellOutput(
            hidden_states=hidden,
            cache=cache,
            state=state,
            attention_weights=attention_weights,
        )