"""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, )