feat: implement sequence-length aware path-c routing and add side channel spec support

Author: apstenku123

README.md

@@ -65,7 +65,8 @@ Repo hygiene:
 - Keep .venv, __pycache__, pytest caches, .beads, agent logs, and
   data/parquet_samples/ out of commits. The Parquet samples are useful for
   local real-data smoke tests, but they are large local artifacts, not repo
-  fixtures.
+  fixtures and remain not checked-in fixtures. The only parquet files allowed
+  in git are committed tiny pytest fixtures under tests/fixtures/parquet/.
 
 Non-goals and limits:
 

cppmega_mlx/data/batch.py

@@ -224,6 +224,31 @@ def as_dict(self, *, include_metadata: bool = False) -> dict[str, Any]:
         return data
 
 
+_DOCUMENT_ID_ALIASES = ("document_ids", "doc_ids", "packing_document_ids")
+
+
+def _document_ids_from_mapping(batch: Mapping[str, Any]) -> Any | None:
+    present = [
+        alias
+        for alias in _DOCUMENT_ID_ALIASES
+        if alias in batch and batch[alias] is not None
+    ]
+    if len(present) > 1:
+        raise ValueError(
+            "only one document-id alias may be provided; got "
+            f"{', '.join(present)}"
+        )
+    if not present:
+        return None
+    alias = present[0]
+    value = batch[alias]
+    has_negative_doc = mx.any(value.astype(mx.int32) < 0)
+    mx.eval(has_negative_doc)
+    if bool(has_negative_doc.item()):
+        raise ValueError(f"{alias} must be non-negative")
+    return value
+
+
 def ensure_lm_batch(batch: LMTokenBatch | Mapping[str, Any] | mx.array) -> LMTokenBatch:
     """Normalize supported tiny-trainer batch inputs into LMTokenBatch."""
 
@@ -239,10 +264,7 @@ def ensure_lm_batch(batch: LMTokenBatch | Mapping[str, Any] | mx.array) -> LMTok
             target_tokens=batch.get("target_tokens", batch.get("target_ids")),
             attention_mask=batch.get("attention_mask"),
             loss_mask=batch.get("loss_mask"),
-            document_ids=batch.get(
-                "document_ids",
-                batch.get("doc_ids", batch.get("packing_document_ids")),
-            ),
+            document_ids=_document_ids_from_mapping(batch),
             structure_ids=batch.get("structure_ids"),
             dep_levels=batch.get("dep_levels"),
             ast_depth_ids=batch.get("ast_depth_ids"),

cppmega_mlx/models/hybrid_lm.py

@@ -837,14 +837,12 @@ def path_c_parameter_logical_aliases(self) -> dict[str, tuple[str, ...]]:
         }
         aliases: dict[str, tuple[str, ...]] = {}
 
-        def add(
-            layer_index: int,
-            parameter_suffix: str,
+        def add_parameter(
+            parameter_name: str,
             logical_suffix: str,
             *,
             block: HybridTinyBlock,
         ) -> None:
-            parameter_name = f"layers.{layer_index}.{parameter_suffix}"
             if parameter_name not in parameter_names:
                 return
             logical_prefixes = [block.path_c_brick_name]
@@ -855,13 +853,27 @@ def add(
                 f"{prefix}_{logical_suffix}" for prefix in logical_prefixes
             )
 
-        for index, block in enumerate(self.layers):
-            add(
-                index,
-                "norm.weight",
+        def add(
+            layer_index: int,
+            parameter_suffix: str,
+            logical_suffix: str,
+            *,
+            block: HybridTinyBlock,
+        ) -> None:
+            add_parameter(
+                f"layers.{layer_index}.{parameter_suffix}",
+                logical_suffix,
+                block=block,
+            )
+
+        for index, block in enumerate(self.layers[:-1]):
+            add_parameter(
+                f"layers.{index + 1}.norm.weight",
                 "residual_norm_weight",
                 block=block,
             )
+
+        for index, block in enumerate(self.layers):
             if block.backend == "mamba3":
                 for parameter_suffix, logical_suffix in (
                     ("block.in_proj.weight", "mamba3_in_proj_weight"),