feat: add side-channel support for training pipelines, fix gradient accumulation in m2rnn TileLang kernels, and improve sidebar UI modularity.

Author: apstenku123

cppmega_mlx/nn/_tilelang/m2rnn_path_c.py

@@ -2658,7 +2658,11 @@ def bwd(
                 skipped = y_val + v_val * d_val
                 dskipped = dpost_val * silu_g
 
-                dy_recurrent[b, t, h, vv] = T.cast(dskipped, grad_dtype)
+                T.atomic_add(
+                    dy_recurrent[b, t, h, vv],
+                    dskipped,
+                    memory_order="relaxed",
+                )
                 T.atomic_add(
                     dconv_input[b, t, v_index],
                     dskipped * d_val,
@@ -2806,7 +2810,11 @@ def bwd(
                     skipped = y_acc[0] + v_val * d_val
                     dskipped = dpost_val * silu_g
 
-                    dy_recurrent[b, t, h, vv] = T.cast(dskipped, grad_dtype)
+                    T.atomic_add(
+                        dy_recurrent[b, t, h, vv],
+                        dskipped,
+                        memory_order="relaxed",
+                    )
                     T.atomic_add(
                         dconv_input[b, t, v_index],
                         dskipped * d_val,
@@ -5384,6 +5392,14 @@ def _apply_vjp(
             v_heads=v_heads,
             g_heads=g_heads,
         )
+        # MLX can otherwise drop the first TileLang owner-output from the
+        # custom-VJP return graph under pytest capture. A scalar readback keeps
+        # the dy buffer live at the autograd boundary without staging or
+        # copying the tensor.
+        try:
+            _ = grads[0][0, 0, 0, 0].item()
+        except IndexError:
+            pass
         return _match_primal_gradient_dtypes(grads, primals)
 
     return _apply

cppmega_mlx/training/compiled.py

@@ -37,7 +37,18 @@
 F = TypeVar("F", bound=Callable[..., Any])
 PathCGradientProbe = Callable[[Mapping[str, Any]], None]
 PathCTrainingRuntime = Any
-PATH_C_TRAINING_VALUE_AND_GRAD_CONTRACT = "path_c_direct_fusion_value_and_grad_v1"
+PATH_C_DIRECT_FUSION_VALUE_AND_GRAD_CONTRACT = (
+    "path_c_direct_fusion_value_and_grad_v1"
+)
+PATH_C_TRAINING_VALUE_AND_GRAD_CONTRACT = (
+    PATH_C_DIRECT_FUSION_VALUE_AND_GRAD_CONTRACT
+)
+PATH_C_FUSED_TRAIN_BLOCK_TRAINING_RUNTIME_CONTRACT = (
+    "path_c_fused_train_block_training_runtime_v1"
+)
+PATH_C_FUSED_TRAIN_BLOCK_VALUE_AND_GRAD_CONTRACT = (
+    "path_c_fused_train_block_value_and_grad_v1"
+)
 
 REGIONAL_COMPILE_TARGETS: Mapping[CompileTarget, bool] = {
     "mamba3_pre": True,
@@ -337,12 +348,16 @@ def attach_path_c_training_runtime(self, runtime: PathCTrainingRuntime) -> None:
         bind = getattr(runtime, "bind_training_graph", None)
         bound = False
         if callable(bind):
-            bind(
-                owner="CompiledPretrainingStep",
-                uses_direct_chain_runtime=True,
-                uses_forward_hook=True,
-                uses_backward_or_vjp_hook=True,
-            )
+            binding = {
+                "owner": "CompiledPretrainingStep",
+                "uses_forward_hook": True,
+                "uses_backward_or_vjp_hook": True,
+            }
+            if _path_c_training_runtime_uses_fused_train_block(runtime):
+                binding["uses_fused_train_block_runtime"] = True
+            else:
+                binding["uses_direct_chain_runtime"] = True
+            bind(**binding)
             bound = True
         try:
             value_and_grad_contract = _path_c_training_runtime_value_and_grad_contract(
@@ -575,7 +590,16 @@ def _path_c_training_runtime_value_and_grad_contract(
     payload = dict(raw_contract)
     contract = str(payload.get("contract", ""))
     owner = str(payload.get("owner", ""))
-    uses_runtime = bool(payload.get("uses_direct_chain_runtime"))
+    uses_direct_chain_runtime = bool(payload.get("uses_direct_chain_runtime"))
+    uses_fused_train_block_runtime = bool(
+        payload.get("uses_fused_train_block_runtime")
+    )
+    direct_contract = contract == PATH_C_DIRECT_FUSION_VALUE_AND_GRAD_CONTRACT
+    fused_contract = contract == PATH_C_FUSED_TRAIN_BLOCK_VALUE_AND_GRAD_CONTRACT
+    uses_runtime = bool(
+        (direct_contract and uses_direct_chain_runtime)
+        or (fused_contract and uses_fused_train_block_runtime)
+    )
     uses_forward = bool(payload.get("uses_forward_hook"))
     uses_reverse = bool(payload.get("uses_backward_or_vjp_hook"))
     returns_model_grads = bool(payload.get("returns_model_grads"))
@@ -586,7 +610,7 @@ def _path_c_training_runtime_value_and_grad_contract(
     hidden_packing = bool(payload.get("hidden_packing_performed", False))
     status = (
         "ok"
-        if contract == PATH_C_TRAINING_VALUE_AND_GRAD_CONTRACT
+        if (direct_contract or fused_contract)
         and owner == "CompiledPretrainingStep"
         and uses_runtime
         and uses_forward
@@ -604,7 +628,8 @@ def _path_c_training_runtime_value_and_grad_contract(
         "status": status,
         "contract": contract or PATH_C_TRAINING_VALUE_AND_GRAD_CONTRACT,
         "owner": owner or None,
-        "uses_direct_chain_runtime": uses_runtime,
+        "uses_direct_chain_runtime": uses_direct_chain_runtime,
+        "uses_fused_train_block_runtime": uses_fused_train_block_runtime,
         "uses_forward_hook": uses_forward,
         "uses_backward_or_vjp_hook": uses_reverse,
         "returns_model_grads": returns_model_grads,
@@ -616,9 +641,22 @@ def _path_c_training_runtime_value_and_grad_contract(
     }
 
 
+def _path_c_training_runtime_uses_fused_train_block(
+    runtime: PathCTrainingRuntime,
+) -> bool:
+    return bool(
+        getattr(runtime, "uses_fused_train_block_runtime", False)
+        or str(getattr(runtime, "contract", ""))
+        == PATH_C_FUSED_TRAIN_BLOCK_TRAINING_RUNTIME_CONTRACT
+    )
+
+
 __all__ = [
     "CompileTarget",
     "CompiledPretrainingStep",
+    "PATH_C_DIRECT_FUSION_VALUE_AND_GRAD_CONTRACT",
+    "PATH_C_FUSED_TRAIN_BLOCK_TRAINING_RUNTIME_CONTRACT",
+    "PATH_C_FUSED_TRAIN_BLOCK_VALUE_AND_GRAD_CONTRACT",
     "PATH_C_TRAINING_VALUE_AND_GRAD_CONTRACT",
     "PathCGradientBufferCapture",
     "PathCGradientProbe",

cppmega_v4/jsonrpc/data_methods.py

@@ -44,11 +44,27 @@ class PreviewRow(BaseModel):
     channels: dict[str, Any] = Field(default_factory=dict)
 
 
+class SideChannelFamilyPreview(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+    family: str
+    status: str
+    columns: list[str] = Field(default_factory=list)
+    missing_columns: list[str] = Field(default_factory=list)
+    dropped_columns: list[str] = Field(default_factory=list)
+    token_alignment: str
+    graph_remapping: str
+    provenance: str
+    non_null_ratio: float
+
+
 class PreviewParquetResult(BaseModel):
     model_config = ConfigDict(extra="forbid")
     rows: list[PreviewRow]
     token_column: str
     available_channels: list[str]
+    side_channel_families: dict[str, SideChannelFamilyPreview] = Field(
+        default_factory=dict
+    )
     bytes_per_token_avg: float
     bytes_per_token_p95: float
     bytes_per_token_max: int
@@ -132,6 +148,20 @@ def preview_parquet(
         rows=rows,
         token_column=token_col,
         available_channels=available,
+        side_channel_families={
+            name: SideChannelFamilyPreview(
+                family=coverage.family,
+                status=coverage.status,
+                columns=list(coverage.columns),
+                missing_columns=list(coverage.missing_columns),
+                dropped_columns=list(coverage.dropped_columns),
+                token_alignment=coverage.token_alignment,
+                graph_remapping=coverage.graph_remapping,
+                provenance=coverage.provenance,
+                non_null_ratio=coverage.non_null_ratio,
+            )
+            for name, coverage in sorted(caps.side_channel_families.items())
+        },
         bytes_per_token_avg=round(bpt_avg, 3),
         bytes_per_token_p95=round(bpt_p95, 3),
         bytes_per_token_max=int(bpt_max),

cppmega_v4/jsonrpc/methods.py

@@ -254,6 +254,34 @@ def _cache_store(cache: LRUCache | None, key: str | None, value: Any):
         cache.set(key, value)
 
 
+def _side_channel_policy_gotchas(
+    side_channels: Any,
+    available: frozenset[str],
+) -> list[GotchaPayload]:
+    """Report GUI-visible contract errors for required side-channel families."""
+    out: list[GotchaPayload] = []
+    global_requires = side_channels.mode == "require"
+    for family, policy in sorted(side_channels.families.items()):
+        if policy.mode == "off":
+            continue
+        required = policy.mode == "require" or global_requires
+        if not required:
+            continue
+        missing = [col for col in policy.columns if col not in available]
+        if not missing:
+            continue
+        out.append(GotchaPayload(
+            id=f"side_channel_required_{family}",
+            severity="error",
+            message=(
+                f"required side-channel family {family!r} is missing "
+                f"{', '.join(missing)}"
+            ),
+            reference="docs/side_channel_conditioning_plan.md#configuration-contract",
+        ))
+    return out
+
+
 # ---------------------------------------------------------------------------
 # verify
 # ---------------------------------------------------------------------------
@@ -339,6 +367,9 @@ def verify(params: VerifyParams, *, cache: LRUCache | None = None) -> VerifyResu
             )
             for g in dverify.gotchas
         ]
+    gotcha_payloads.extend(
+        _side_channel_policy_gotchas(params.side_channels, available)
+    )
 
     edge_payloads: list[EdgeResolution] = []
     for re in resolved.edges:

cppmega_v4/runner/stages.py

@@ -424,6 +424,16 @@ def loss_fn(model: nn.Module, emb: mx.array, tgt: mx.array) -> mx.array:
         data_source = "synthetic"
         token_count = 0
         tokenizer_used: str | None = None
+        # V4-10: record which side-channels the caller supplied. Pure
+        # observation for now — actual per-channel forward routing is
+        # v5+ work. Surfacing via extras lets e2e prove UI toggles
+        # reached the backend opts surface.
+        side_channels_in = opts.get("side_channels") or {}
+        side_channels_observed: list[str] = []
+        if isinstance(side_channels_in, dict):
+            for name, data in side_channels_in.items():
+                if isinstance(data, (list, tuple)) and len(data) > 0:
+                    side_channels_observed.append(str(name))
         parquet_path = opts.get("parquet_path")
         tokenizer_path = opts.get("tokenizer_path")
         targets = mx.random.randint(0, vocab_size, shape=(batch, seq))
@@ -608,6 +618,7 @@ def _count(tree: Any) -> int:
                     "l2_diff": round(l2_diff, 6),
                     "cos_sim": round(cos_sim, 6),
                 },
+                "side_channels_observed": side_channels_observed,
                 "model_summary": _summarize_model(
                     ctx.spec, optimizer_kind, schedule_kind_label),
             },