feat: add support for backward stage indexing and refine path-c ABI output management and ELK UI snake layout logic

Author: apstenku123

cppmega_mlx/runtime/path_c_fusion_schedules.py

@@ -56,6 +56,7 @@
     "path_c_run_backward": 1,
     "path_c_row_chunk_index": 0,
     "path_c_row_subchunk_index": 0,
+    "path_c_backward_stage_index": 0,
 }
 
 REGIONAL_COMPILE_TARGETS: Mapping[CompileTarget, bool] = {
@@ -148,6 +149,8 @@ def __init__(
         row_subchunk_count: int | None = None,
         row_subchunk_index_param: str | None = None,
         rows_per_kernel_launch: int | None = None,
+        backward_stage_count: int | None = None,
+        backward_stage_index_param: str | None = None,
     ) -> None:
         if not callable(kernel):
             raise TypeError("fused train-block kernel must be callable")
@@ -200,12 +203,24 @@ def __init__(
             if rows_per_kernel_launch is not None
             else None
         )
+        self.backward_stage_count = (
+            max(1, int(backward_stage_count))
+            if backward_stage_count is not None
+            else None
+        )
+        self.backward_stage_index_param = (
+            str(backward_stage_index_param) if backward_stage_index_param else None
+        )
         self._cppmega_path_c_backward_gate_param = self.backward_gate_param
         self._cppmega_path_c_row_chunk_count = self.row_chunk_count
         self._cppmega_path_c_row_chunk_index_param = self.row_chunk_index_param
         self._cppmega_path_c_row_subchunk_count = self.row_subchunk_count
         self._cppmega_path_c_row_subchunk_index_param = self.row_subchunk_index_param
         self._cppmega_path_c_rows_per_kernel_launch = self.rows_per_kernel_launch
+        self._cppmega_path_c_backward_stage_count = self.backward_stage_count
+        self._cppmega_path_c_backward_stage_index_param = (
+            self.backward_stage_index_param
+        )
         self._logical_gradient_names = frozenset(
             name for name in self.physical_abi_map if name.endswith("_grad")
         )
@@ -486,6 +501,8 @@ def value_and_grad_contract(self) -> Mapping[str, Any]:
             "row_subchunk_count": self.row_subchunk_count,
             "row_subchunk_index_param": self.row_subchunk_index_param,
             "rows_per_kernel_launch": self.rows_per_kernel_launch,
+            "backward_stage_count": self.backward_stage_count,
+            "backward_stage_index_param": self.backward_stage_index_param,
         }
 
 

cppmega_mlx/training/compiled.py

@@ -125,6 +125,8 @@ def build_fused_replay_boundary_custom_function(
     row_chunk_index_param: str | None = None,
     row_subchunk_count: int | None = None,
     row_subchunk_index_param: str | None = None,
+    backward_stage_count: int | None = None,
+    backward_stage_index_param: str | None = None,
     backward_gate_param: str = "path_c_run_backward",
 ) -> FusedReplayBoundaryCallable:
     """Return ``f(hidden_entry, *params) -> boundary outputs``.
@@ -160,6 +162,8 @@ def build_fused_replay_boundary_custom_function(
         max(1, int(row_subchunk_count or 1)) if subchunk_param else 1
     )
     gate_param = str(backward_gate_param or "path_c_run_backward")
+    stage_param = str(backward_stage_index_param or "")
+    stage_count = max(1, int(backward_stage_count or 1)) if stage_param else 1
     backward_sync_outputs = (
         f"{hidden_entry_logical_name}_grad",
         *(
@@ -192,10 +196,21 @@ def _launch_all(*, run_backward: bool) -> None:
             row_chunk_index_param=chunk_param,
             row_subchunk_index_param=subchunk_param,
         ):
-            artifact.forward(
-                bank_owner=bank_owner,
-                kernel_scalar_params=scalar_params,
-            )
+            stage_range = range(stage_count) if run_backward else range(1)
+            for stage_index in stage_range:
+                stage_scalar_params = dict(scalar_params)
+                if run_backward and stage_param:
+                    stage_scalar_params[stage_param] = int(stage_index)
+                try:
+                    artifact.forward(
+                        bank_owner=bank_owner,
+                        kernel_scalar_params=stage_scalar_params,
+                    )
+                except Exception as exc:
+                    raise RuntimeError(
+                        "fused replay artifact launch failed with "
+                        f"kernel_scalar_params={stage_scalar_params!r}"
+                    ) from exc
         sync_outputs = backward_sync_outputs if run_backward else boundary_outputs
         mx.eval(*_logical_views(abi_map, bank_buffers, sync_outputs))
 

cppmega_mlx/training/path_c_fused_replay.py

@@ -1,14 +1,14 @@
-bundle: http://127.0.0.1:8765/assets/index-B0DK3qBx.js
+bundle: http://127.0.0.1:8765/assets/index-DoLZxAnf.js
 step 1 action: load preset deepseek_v4_flash
 step 2 action: Auto Align Graph (folded)
   [folded] 9 edges, crossings (interior) = 0
 step 3 action: Unpack All, then Auto Align Graph
   [unpacked] 62 edges, crossings (interior) = 6
-    X  engram_unpacked_r0_1779964444764_2_add -> dsv4_attentio  ×  dsv4_attention_unpacked_r0_1779964444764_3 -> dsv4_atte
-    X  engram_unpacked_r0_1779964444764_2_add -> dsv4_attentio  ×  engram_unpacked_r0_1779964444764_2_add -> dsv4_attentio
-    X  csa_hca_unpacked_r1_1779964444764_1_add -> engram_unpac  ×  engram_unpacked_r1_1779964444764_2 -> engram_unpacked_r
-    X  csa_hca_unpacked_r1_1779964444764_1_add -> engram_unpac  ×  csa_hca_unpacked_r1_1779964444764_1_add -> engram_unpac
-    X  lightning_indexer_unpacked_r2_1779964444764_0_add -> cs  ×  csa_hca_unpacked_r2_1779964444764_1 -> csa_hca_unpacked
-    X  lightning_indexer_unpacked_r2_1779964444764_0_add -> cs  ×  lightning_indexer_unpacked_r2_1779964444764_0_add -> cs
+    X  engram_unpacked_r0_1779964683031_2_add -> dsv4_attentio  ×  dsv4_attention_unpacked_r0_1779964683031_3 -> dsv4_atte
+    X  engram_unpacked_r0_1779964683031_2_add -> dsv4_attentio  ×  engram_unpacked_r0_1779964683031_2_add -> dsv4_attentio
+    X  csa_hca_unpacked_r1_1779964683031_1_add -> engram_unpac  ×  engram_unpacked_r1_1779964683031_2 -> engram_unpacked_r
+    X  csa_hca_unpacked_r1_1779964683031_1_add -> engram_unpac  ×  csa_hca_unpacked_r1_1779964683031_1_add -> engram_unpac
+    X  lightning_indexer_unpacked_r2_1779964683031_0_add -> cs  ×  csa_hca_unpacked_r2_1779964683031_1 -> csa_hca_unpacked
+    X  lightning_indexer_unpacked_r2_1779964683031_0_add -> cs  ×  lightning_indexer_unpacked_r2_1779964683031_0_add -> cs
 RESULT folded   crossings = 0
 RESULT unpacked crossings = 6