Add MoE workspace and diagnostics

src/initium/LNN_controller/cfc_backbone_parallel.py

@@ -100,6 +100,8 @@ def __init__(
         moe_top_k: int = 1,
         moe_capacity_factor: float = 1.5,
         moe_load_balance_alpha: float = 0.01,
+        moe_residual_scale: float = 1.0,
+        moe_block_isolation: int | None = None,
         device=None,
         dtype=None,
     ):
@@ -134,32 +136,68 @@ def __init__(
         # token batch, so this costs nothing extra to support here).
         self.moe_enabled = moe_enabled
         self.moe_blocks: nn.ModuleList | None = None
+        if moe_block_isolation is not None and not 0 <= moe_block_isolation < num_blocks:
+            raise ValueError(
+                f"moe_block_isolation={moe_block_isolation} is outside "
+                f"the configured backbone block range [0, {num_blocks - 1}]"
+            )
         if moe_enabled:
             self.moe_blocks = nn.ModuleList(
                 [
-                    MoEBlock(
-                        d_model,
-                        num_experts=moe_num_experts,
-                        expert_dim=moe_expert_dim,
-                        top_k=moe_top_k,
-                        capacity_factor=moe_capacity_factor,
-                        load_balance_alpha=moe_load_balance_alpha,
-                        device=device,
-                        dtype=dtype,
+                    (
+                        MoEBlock(
+                            d_model,
+                            num_experts=moe_num_experts,
+                            expert_dim=moe_expert_dim,
+                            top_k=moe_top_k,
+                            capacity_factor=moe_capacity_factor,
+                            load_balance_alpha=moe_load_balance_alpha,
+                            residual_scale=moe_residual_scale,
+                            device=device,
+                            dtype=dtype,
+                        )
+                        if moe_block_isolation is None or i == moe_block_isolation
+                        else nn.Identity()
                     )
-                    for _ in range(num_blocks)
+                    for i in range(num_blocks)
                 ]
             )
 
-    def forward(self, x):  # (B, L, in_dim) -> (B, L, d_model), called ONCE per training step
+    def forward(
+        self, x, *, token_mask: torch.Tensor | None = None
+    ):  # (B, L, in_dim) -> (B, L, d_model), called ONCE per training step
         h = self.in_adapter(x)
         for i, block in enumerate(self.blocks):
             h = block(h)
-            if self.moe_enabled:
+            if self.moe_enabled and isinstance(self.moe_blocks[i], MoEBlock):
                 assert self.moe_blocks is not None
-                h = self.moe_blocks[i](h)
+                h = self.moe_blocks[i](h, token_mask=token_mask)
         return h
 
+    def forward_with_specialists(
+        self, x: torch.Tensor, *, token_mask: torch.Tensor | None = None
+    ):
+        """Return final MoE candidates and routes for shared-workspace use."""
+        if not self.moe_enabled or self.moe_blocks is None:
+            raise RuntimeError("workspace broadcast requires an enabled MoE expert pool")
+        h = self.in_adapter(x)
+        final_index = len(self.blocks) - 1
+        if not isinstance(self.moe_blocks[final_index], MoEBlock):
+            raise RuntimeError(
+                "workspace broadcast requires the final backbone block to have an active MoE bank"
+            )
+        for i, block in enumerate(self.blocks):
+            h = block(h)
+            if i == final_index:
+                base = h
+                _, experts, indices, weights = self.moe_blocks[i](
+                    h, return_details=True, token_mask=token_mask
+                )
+                return base, experts, indices, weights
+            if isinstance(self.moe_blocks[i], MoEBlock):
+                h = self.moe_blocks[i](h, token_mask=token_mask)
+        raise RuntimeError("CfC parallel backbone has no blocks")
+
 
 def build_parallel_backbone(
     in_dim,
@@ -177,6 +215,8 @@ def build_parallel_backbone(
     moe_top_k: int = 1,
     moe_capacity_factor: float = 1.5,
     moe_load_balance_alpha: float = 0.01,
+    moe_residual_scale: float = 1.0,
+    moe_block_isolation: int | None = None,
     device=None,
     dtype=None,
 ):
@@ -204,6 +244,8 @@ def build_parallel_backbone(
                     moe_top_k=moe_top_k,
                     moe_capacity_factor=moe_capacity_factor,
                     moe_load_balance_alpha=moe_load_balance_alpha,
+                    moe_residual_scale=moe_residual_scale,
+                    moe_block_isolation=moe_block_isolation,
                     **(cfc_kwargs or {}),
                 )
             )
@@ -226,6 +268,8 @@ def build_parallel_backbone(
                     moe_top_k=moe_top_k,
                     moe_capacity_factor=moe_capacity_factor,
                     moe_load_balance_alpha=moe_load_balance_alpha,
+                    moe_residual_scale=moe_residual_scale,
+                    moe_block_isolation=moe_block_isolation,
                 )
             )
         cur = d_model
@@ -246,7 +290,9 @@ def collect_backbone_moe_layers(backbone: nn.Module) -> list:
                 raise TypeError(
                     "MoE-enabled backbone stages must expose a ModuleList of moe_blocks"
                 )
-            moe_layers.extend(blocks)
+            moe_layers.extend(
+                block for block in blocks if callable(getattr(block, "pop_aux_loss", None))
+            )
     return moe_layers
 
 

src/initium/LNN_controller/cfc_controller.py

@@ -113,6 +113,15 @@ def step(self, x, state):
         y = x + out.to(x.dtype) if self.residual else out.to(x.dtype)
         return y, _state_to_fp32(new_state)
 
+    def forward(self, x, state):
+        """Run one recurrent step through ``Module.__call__``.
+
+        Keeping the step behind ``forward`` preserves the equations while
+        ensuring module hooks (including segment activation diagnostics) see
+        the CfC block boundary.
+        """
+        return self.step(x, state)
+
 
 class CfCControllerWrapper(nn.Module):
     """Stack of `num_blocks` CfCControllerBlocks with the nn.LSTM-compatible
@@ -247,6 +256,12 @@ def init_state(self, batch_size, device=None, dtype=None):
         return [blk.init_state(batch_size, device=device) for blk in blocks]
 
     def forward(self, input, hx):
+        if isinstance(input, list):
+            if self.moe_source_dims is None:
+                raise RuntimeError(
+                    "list input is supported only when moe_source_dims is configured"
+                )
+            return self.forward_multi_source(input, hx)
         assert input.dim() == 3 and input.size(1) == 1, (
             "CfCControllerWrapper only supports single-timestep calls "
             f"(got shape {tuple(input.shape)})."
@@ -262,7 +277,7 @@ def forward(self, input, hx):
         new_hx = []
         blocks = cast(Iterable[CfCControllerBlock], self.blocks)
         for i, (block, state) in enumerate(zip(blocks, hx)):
-            x, new_state = block.step(x, state)
+            x, new_state = block(x, state)
             if self.moe_enabled:
                 assert self.moe_blocks is not None
                 x = self.moe_blocks[i](x)
@@ -283,15 +298,27 @@ def forward_multi_source(self, sources: list[torch.Tensor], hx):
         if self.moe_source_dims is None:
             raise RuntimeError("forward_multi_source requires moe_source_dims to have been set")
         assert self.source_in_adapters is not None and self.moe_blocks is not None
+        # Preserve the stock CfC input-fusion path exactly: the baseline
+        # receives one affine projection of concat([backbone, previous
+        # read]). The source-aware MoE is an additive residual around that
+        # fused representation. This keeps the memory signal and controller
+        # initialization intact while still allowing the experts to learn
+        # source-specific corrections.
+        concatenated = torch.cat(sources, dim=-1)
+        x = self.in_adapter(concatenated)
         projected = [adapter(s) for adapter, s in zip(self.source_in_adapters, sources)]
-        fused_per_source = self.moe_blocks[0](projected)  # list[Tensor], one per source
-        x = torch.stack(fused_per_source, dim=0).sum(dim=0)  # combine into one fused input
+        fused_per_source = self.moe_blocks[0](projected)  # source + routed delta
+        moe_delta = torch.stack(
+            [fused - source for fused, source in zip(fused_per_source, projected)],
+            dim=0,
+        ).sum(dim=0)
+        x = x + moe_delta
         if hx is None:
             hx = self.init_state(x.size(0), device=x.device)
         new_hx = []
         blocks = cast(Iterable[CfCControllerBlock], self.blocks)
         for block, state in zip(blocks, hx):
-            x, new_state = block.step(x, state)
+            x, new_state = block(x, state)
             new_hx.append(new_state)
         return x.unsqueeze(1), new_hx