Path C: fail closed on unverified bank-grad overlay

Close the correctness holes found in the parameter-bank residency follow-up.

Root cause 1: bank residency alone is not enough to consume fused gradients.
The fused train-block needs runtime inputs (hidden_entry, target_ids,
target_mask) written into the ABI banks before artifact.value_and_grad runs.
Without the suffix-bypass custom-function surface, merged mode ran the kernel
against stale/zero inputs and then overwrote correct eager in-region grads
with bogus bank values. PathCFusedPlusEagerTrainingRuntime now fails closed:
when suffix-bypass is absent, it keeps the old warmup+eager behaviour and
never calls artifact.value_and_grad or overlays bank grad slots.

Root cause 2: the current fused suffix ABI advertises pending loss codegen.
The m04 installer now refuses to attach the suffix-bypass loss function when
_cppmega_path_c_train_step_suffix_loss_input_abi.reason contains pending, or
when the scalar output ABI does not report computed outputs. This keeps the
route available while preventing unverified suffix numerics from entering the
training critical path.

Also tighten write_into_bank_slot: dtype must match the generated ABI dtype
before slice assignment. MLX silently coerces float32 into int32 bank slots;
this now raises instead of hiding a dtype/layout bug.

Verification:
- ruff on touched files
- pyright on cppmega_mlx/runtime/path_c_physical_abi.py, cppmega_mlx/training/compiled.py, cppmega_mlx/training/path_c_fused_suffix.py; checked no new pyright errors in the touched m04 factory window
- pytest tests/test_path_c_physical_abi.py tests/test_path_c_fused_plus_eager_runtime.py tests/test_path_c_fused_suffix_custom_function.py tests/test_hybrid_lm_path_c_physical_abi_bank_owner.py tests/test_compiled_train.py -q: 91 passed
- pytest tests/test_m04_train_step.py -k "path_c or fused_train_block" -q: 66 passed
- pytest tests/v4/test_fusion_stage_{a,b,c,d,e,f}.py tests/v4/test_fusion_roadmap_gaps.py -q: 401 passed
- live local_gb10_quarter route: status=m04_path_c_training_route_available, suffix_bypass_available=False, bank_grad_overlay_active=False, merged_bank_resident_parameter_count=0, layers.10.block.D eager grad norm non-zero

Remaining work: re-enable suffix-bypass only after the TileLang generated suffix
loss/gradient code passes numerical parity against the eager path. Until then
Path C stays honest warmup/eager instead of silently training the wrong model.

Author: apstenku123

cppmega_mlx/runtime/path_c_physical_abi.py

@@ -633,6 +633,13 @@ def write_into_bank_slot(
     bank = bank_buffers[bank_name]
     offset = int(info.get("offset", 0) or 0)
     size = int(info.get("size", 1) or 1)
+    expected_dtype = str(info.get("dtype", ""))
+    actual_dtype = _dtype_of(value)
+    if expected_dtype and actual_dtype != expected_dtype:
+        raise ValueError(
+            f"value for {logical_name!r} has dtype {actual_dtype!r}, "
+            f"expected {expected_dtype!r}"
+        )
     # Flatten the value into a 1-D array of length ``size``.
     expected_size = int(prod(tuple(int(dim) for dim in tuple(getattr(value, "shape", ())))))
     if expected_size != size:

cppmega_mlx/training/compiled.py

@@ -710,7 +710,15 @@ def value_and_grad(
                 self.last_fused_value_and_grad_payload = vg_payload
                 raise
 
-        merged_mode = bool(self.in_region_parameter_bank_aliases)
+        # Bank residency alone is not enough to consume fused gradients:
+        # the fused train-block also needs per-batch runtime inputs
+        # (hidden_entry, target_ids, target_mask) written into the ABI banks.
+        # Without the suffix-bypass custom-function surface, the artifact's
+        # value_and_grad would run against stale/zero runtime inputs and would
+        # silently overwrite correct eager grads with bogus bank values. Fail
+        # closed to the original warmup+eager behavior unless suffix-bypass
+        # handled the step above.
+        merged_mode = False
 
         # In merged mode, push optimizer-replaced parameters into bank slots
         # before the fused kernel reads them. The first step's sync is a no-op
@@ -745,6 +753,7 @@ def value_and_grad(
             "missing_parameter_names": (),
             "bank_sync_status": bank_sync_status,
             "error": None,
+            "suffix_bypass_active": False,
         }
 
         fused_grads: Any | None = None
@@ -832,19 +841,40 @@ def value_and_grad_contract(self) -> Mapping[str, Any]:
         inner = dict(self.artifact.value_and_grad_contract())
         inner_coverage = inner.get("full_model_gradient_coverage")
         coverage = dict(inner_coverage) if isinstance(inner_coverage, Mapping) else {}
-        merged_mode_active = bool(self.in_region_parameter_bank_aliases)
+        bank_residency_ready = bool(self.in_region_parameter_bank_aliases)
+        suffix_bypass_available = bool(
+            self.fused_suffix_loss_fn is not None
+            and self.in_region_parameter_bank_aliases
+        )
         reason = (
-            "mixed-mode training runtime returns full-model grads via merged "
-            "bank-resident fused grads + eager residual grads: in-region "
-            "parameters are bank views and the fused artifact populates "
-            "their gradients into the same bank storage"
-            if merged_mode_active
+            "mixed-mode training runtime returns full-model grads via "
+            "suffix-bypass: in-region parameters and hidden-entry cotangents "
+            "come from the fused custom-function VJP, residual/prefix grads "
+            "come from MLX eager autograd"
+            if suffix_bypass_available
             else (
-                "mixed-mode training runtime returns full-model grads: "
-                "fused in-region warmup is observable, residual grads "
-                "come from the trainer's eager value_and_grad closure"
+                "mixed-mode training runtime returns full-model grads from "
+                "the trainer's eager value_and_grad closure; bank-resident "
+                "fused gradients are disabled because no verified runtime "
+                "input writer / suffix-bypass bridge is active"
+                if bank_residency_ready
+                else (
+                    "mixed-mode training runtime returns full-model grads: "
+                    "fused in-region warmup is observable, residual grads "
+                    "come from the trainer's eager value_and_grad closure"
+                )
             )
         )
+        overlay_parameter_count = (
+            len(self.in_region_parameter_bank_aliases)
+            if suffix_bypass_available
+            else 0
+        )
+        overlay_parameter_names = (
+            tuple(sorted(self.in_region_parameter_bank_aliases))
+            if suffix_bypass_available
+            else ()
+        )
         coverage.update(
             {
                 "full_model_gradient_tree_ready": True,
@@ -858,17 +888,10 @@ def value_and_grad_contract(self) -> Mapping[str, Any]:
                 "fused_in_region_parameter_count": len(
                     self.in_region_parameter_names
                 ),
-                "parameter_bank_residency_active": merged_mode_active,
-                "merged_bank_resident_parameter_count": (
-                    len(self.in_region_parameter_bank_aliases)
-                    if merged_mode_active
-                    else 0
-                ),
-                "merged_bank_resident_parameter_names": (
-                    tuple(sorted(self.in_region_parameter_bank_aliases))
-                    if merged_mode_active
-                    else ()
-                ),
+                "parameter_bank_residency_active": bank_residency_ready,
+                "bank_grad_overlay_active": suffix_bypass_available,
+                "merged_bank_resident_parameter_count": overlay_parameter_count,
+                "merged_bank_resident_parameter_names": overlay_parameter_names,
             }
         )
         merged = {
@@ -893,13 +916,9 @@ def value_and_grad_contract(self) -> Mapping[str, Any]:
             "parameter_bank_residency_active": bool(
                 self.in_region_parameter_bank_aliases
             ),
-            "merged_bank_resident_parameter_count": len(
-                self.in_region_parameter_bank_aliases
-            ),
-            "suffix_bypass_available": bool(
-                self.fused_suffix_loss_fn is not None
-                and self.in_region_parameter_bank_aliases
-            ),
+            "bank_grad_overlay_active": suffix_bypass_available,
+            "merged_bank_resident_parameter_count": overlay_parameter_count,
+            "suffix_bypass_available": suffix_bypass_available,
             "runtime_class": type(self).__name__,
         }
         return merged

scripts/m04_train_step.py

@@ -3109,6 +3109,25 @@ def _build_path_c_fused_suffix_loss_fn_for_model(
     required_inputs = ("target_ids", "target_mask", "loss", "ntokens")
     if not all(name in abi_map for name in required_inputs):
         return None
+    suffix_input_abi = getattr(
+        prim_func,
+        "_cppmega_path_c_train_step_suffix_loss_input_abi",
+        {},
+    )
+    suffix_input_reason = ""
+    if isinstance(suffix_input_abi, Mapping):
+        suffix_input_reason = str(suffix_input_abi.get("reason", ""))
+    if "pending" in suffix_input_reason.lower():
+        return None
+    output_abi = getattr(
+        prim_func,
+        "_cppmega_path_c_train_step_output_abi",
+        {},
+    )
+    if isinstance(output_abi, Mapping) and not bool(
+        output_abi.get("outputs_computed", False)
+    ):
+        return None
 
     first_in_region_layer_index_fn = getattr(
         model, "path_c_fused_first_in_region_layer_index", None
@@ -3120,6 +3139,7 @@ def _build_path_c_fused_suffix_loss_fn_for_model(
     )
     if first_in_region_layer_index is None:
         return None
+    first_in_region_layer_index_int = int(cast(int, first_in_region_layer_index))
 
     parameter_order = tuple(sorted(in_region_parameter_bank_aliases.keys()))
     try:
@@ -3140,7 +3160,7 @@ def _build_path_c_fused_suffix_loss_fn_for_model(
     model.attach_path_c_fused_suffix_custom_function(
         fused_suffix,
         parameter_order=parameter_order,
-        first_in_region_layer_index=int(first_in_region_layer_index),
+        first_in_region_layer_index=first_in_region_layer_index_int,
     )
 
     def _loss_fn(loss_model: Any, batch: Mapping[str, Any]) -> tuple[mx.array, mx.array]:

tests/test_m04_train_step.py

@@ -4809,6 +4809,10 @@ def fake_lowerer(func_or_mod: Any, *, target: str, **kwargs: Any) -> Any:
     assert coverage["full_model_gradient_tree_ready"] is True
     assert coverage["gradient_scope"] == "full_model_via_mixed_mode"
     assert coverage["missing_parameter_count"] == 0
+    assert value_and_grad_contract["suffix_bypass_available"] is False
+    assert value_and_grad_contract["bank_grad_overlay_active"] is False
+    assert value_and_grad_contract["merged_bank_resident_parameter_count"] == 0
+    assert coverage["bank_grad_overlay_active"] is False
     assert value_and_grad_contract["delegates_to_eager_loss_and_grad"] is False
     assert value_and_grad_contract["hidden_packing_performed"] is False
     inner = model.path_c_fused_train_block_artifact.value_and_grad_contract()

tests/test_path_c_fused_plus_eager_runtime.py

@@ -565,7 +565,7 @@ def __init__(self) -> None:
         }
 
 
-def test_merged_mode_overlays_bank_resident_grads_onto_eager_tree() -> None:
+def test_bank_residency_without_suffix_bypass_fails_closed_to_eager_grads() -> None:
     art = _FakeFusedArtifactWithBankResidentGrads()
     # Two in-region params, each with a unique sentinel grad value the runtime
     # must surface to the trainer (replacing whatever the eager closure
@@ -631,39 +631,34 @@ def closure(model: Any, batch: Mapping[str, Any]) -> Any:
         batch={},
         loss_and_grad=closure,
     )
-    # Sync callable was invoked exactly once.
-    assert len(sync_calls) == 1
-    # The fused artifact value_and_grad was invoked with bank_owner threaded in.
-    assert len(art.value_and_grad_calls) == 1
-    assert (
-        art.value_and_grad_calls[-1]["kwargs"].get("bank_owner") is bank_owner
-    )
-    # In-region grads now come from the fused tree (sentinel values), not the
-    # eager closure (which produced 99s).
+    # Without suffix-bypass / runtime-input writing, the runtime must not sync
+    # params into banks, must not call artifact.value_and_grad, and must not
+    # overlay fused grad slots. That path would run against stale hidden/target
+    # bank inputs and corrupt the eager gradients.
+    assert len(sync_calls) == 0
+    assert len(art.value_and_grad_calls) == 0
+    assert len(art.forward_calls) == 1
     flat_grads = dict(tree_flatten(grads))
-    assert flat_grads["layers.10.block.D"].tolist() == [1.0, 2.0, 3.0, 4.0]
-    assert flat_grads["layers.11.block.D"].tolist() == [5.0, 6.0, 7.0, 8.0]
+    assert flat_grads["layers.10.block.D"].tolist() == [99.0, 99.0, 99.0, 99.0]
+    assert flat_grads["layers.11.block.D"].tolist() == [99.0, 99.0, 99.0, 99.0]
     # Out-of-region grads are unchanged (eager wins).
     assert flat_grads["layers.0.block.q_proj.weight"].shape == (16, 16)
     assert flat_grads["lm_head.weight"].shape == (256, 16)
-    # Telemetry reflects merged-mode behaviour.
-    assert runtime.last_fused_warmup_payload["attempted"] is False
+    # Telemetry reflects fail-closed warmup+eager behaviour.
+    assert runtime.last_fused_warmup_payload["attempted"] is True
     vg = runtime.last_fused_value_and_grad_payload
-    assert vg["attempted"] is True
-    assert vg["completed"] is True
-    assert vg["elapsed_ns"] >= 0
-    assert vg["merged_parameter_count"] == 2
-    assert vg["merged_parameter_names"] == (
-        "layers.10.block.D",
-        "layers.11.block.D",
-    )
+    assert vg["attempted"] is False
+    assert vg["completed"] is False
+    assert vg["merged_parameter_count"] == 0
+    assert vg["merged_parameter_names"] == ()
     assert vg["missing_parameter_names"] == ()
-    assert vg["bank_sync_status"] == {"status": "ok", "synced": ["x"], "skipped": []}
+    assert vg["bank_sync_status"] is None
+    assert vg["suffix_bypass_active"] is False
 
 
-def test_merged_mode_records_missing_bank_grads_without_dropping_coverage() -> None:
-    """If the fused tree is missing a bank-bound parameter, the merger keeps
-    the eager grad in place and records the missing parameter for telemetry."""
+def test_bank_residency_without_suffix_bypass_never_reads_partial_bank_grads() -> None:
+    """Partial fused grad trees are ignored until suffix-bypass supplies
+    verified runtime inputs for the kernel."""
 
     art = _FakeFusedArtifactWithBankResidentGrads()
     # Only one of the two in-region params has a bank-resident grad.
@@ -696,14 +691,15 @@ def closure(model: Any, batch: Mapping[str, Any]) -> Any:
         loss_and_grad=closure,
     )
     flat_grads = dict(tree_flatten(grads))
-    # Bank-resident grad replaced the eager one.
-    assert flat_grads["layers.10.block.D"].tolist() == [1.0, 2.0, 3.0, 4.0]
-    # Missing-from-fused grad kept the eager value.
+    # Even the available bank-resident grad is ignored in fail-closed mode;
+    # eager remains the source of truth for all parameters.
+    assert flat_grads["layers.10.block.D"].tolist() == [0.0, 0.0, 0.0, 0.0]
     assert flat_grads["layers.11.block.D"].tolist() == [42.0, 42.0, 42.0, 42.0]
     vg = runtime.last_fused_value_and_grad_payload
-    assert vg["merged_parameter_count"] == 1
-    assert vg["merged_parameter_names"] == ("layers.10.block.D",)
-    assert vg["missing_parameter_names"] == ("layers.11.block.D",)
+    assert vg["merged_parameter_count"] == 0
+    assert vg["merged_parameter_names"] == ()
+    assert vg["missing_parameter_names"] == ()
+    assert vg["suffix_bypass_active"] is False
 
 
 def test_merged_mode_contract_exposes_bank_residency_signals() -> None:
@@ -719,16 +715,14 @@ def test_merged_mode_contract_exposes_bank_residency_signals() -> None:
     )
     contract = runtime.value_and_grad_contract()
     assert contract["parameter_bank_residency_active"] is True
-    assert contract["merged_bank_resident_parameter_count"] == 2
+    assert contract["bank_grad_overlay_active"] is False
+    assert contract["merged_bank_resident_parameter_count"] == 0
     coverage = contract["full_model_gradient_coverage"]
     assert coverage["parameter_bank_residency_active"] is True
-    assert coverage["merged_bank_resident_parameter_count"] == 2
-    assert coverage["merged_bank_resident_parameter_names"] == (
-        "layers.10.block.D",
-        "layers.11.block.D",
-    )
-    # Reason includes the merged-mode phrasing.
-    assert "bank-resident" in coverage["reason"]
+    assert coverage["bank_grad_overlay_active"] is False
+    assert coverage["merged_bank_resident_parameter_count"] == 0
+    assert coverage["merged_bank_resident_parameter_names"] == ()
+    assert "disabled" in coverage["reason"]
 
 
 def test_warmup_mode_contract_clears_bank_residency_signals() -> None:
@@ -741,9 +735,11 @@ def test_warmup_mode_contract_clears_bank_residency_signals() -> None:
     )
     contract = runtime.value_and_grad_contract()
     assert contract["parameter_bank_residency_active"] is False
+    assert contract["bank_grad_overlay_active"] is False
     assert contract["merged_bank_resident_parameter_count"] == 0
     coverage = contract["full_model_gradient_coverage"]
     assert coverage["parameter_bank_residency_active"] is False
+    assert coverage["bank_grad_overlay_active"] is False
     assert coverage["merged_bank_resident_parameter_count"] == 0
     assert coverage["merged_bank_resident_parameter_names"] == ()
 
@@ -777,6 +773,7 @@ def _stub_loss(model: Any, batch: Mapping[str, Any]) -> tuple[mx.array, mx.array
         runtime_with_bypass.value_and_grad_contract()["suffix_bypass_available"]
         is True
     )
+    assert runtime_with_bypass.value_and_grad_contract()["bank_grad_overlay_active"] is True
 
 
 def test_suffix_bypass_uses_fused_suffix_loss_fn_and_skips_eager_closure() -> None:

tests/test_path_c_physical_abi.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import pytest
+import mlx.core as mx
 
 from cppmega_mlx.runtime.path_c_physical_abi import (
     PathCLogicalBufferOwner,
@@ -13,6 +14,7 @@
     plan_physical_abi_runtime_bridge,
     validate_physical_abi_map,
     validate_physical_abi_runtime_bindings,
+    write_into_bank_slot,
 )
 
 
@@ -436,3 +438,25 @@ def test_runtime_binding_rejects_logical_tensor_substitutes_for_banked_abi() ->
     assert payload["missing_bank_buffers"] == ["path_c_float32_abi_bank"]
     assert payload["unexpected_buffers"] == ["hidden", "out"]
     assert any("missing caller-owned bank buffer" in error for error in payload["errors"])
+
+
+def test_write_into_bank_slot_rejects_dtype_mismatch() -> None:
+    mapping = {
+        "target_ids": {
+            "bank": "path_c_int32_abi_bank",
+            "dtype": "int32",
+            "offset": 0,
+            "shape": (4,),
+            "logical_shape": (4,),
+            "size": 4,
+        },
+    }
+    buffers = {"path_c_int32_abi_bank": mx.zeros((4,), dtype=mx.int32)}
+
+    with pytest.raises(ValueError, match="dtype 'float32'.*expected 'int32'"):
+        write_into_bank_slot(
+            mapping,
+            buffers,
+            "target_ids",
+            mx.array([1.5, 2.5, 3.5, 4.5], dtype=mx.float32),
+        )