fix(path-c): time-chunk m2rnn_bwd too (not just mamba3) so direct-chain backward clears the Metal watchdog

The intra-segment time-chunking added in ac412fb switched ONLY the
mamba3_mimo_bwd reverse-time-scan backward segment to launcher_chunks
dispatch to stay under the macOS GPU watchdog. But m2rnn_bwd is an
equally long reverse-time recurrent scan and, in the direct-chain reverse
order, its segment runs BEFORE mamba3. Under the default grid_chunks
dispatch it ran as one monolithic Metal command buffer and tripped the
watchdog FIRST -> kIOGPUCommandBufferCallbackErrorTimeout at the
m2rnn_bwd segment, before mamba3 was ever reached.

Verified by running the real m04 --use-path-c-direct-chain-runtime route
on Metal: the direct-chain training runtime (PathCDirectFusionChainTraining
Runtime.value_and_grad) IS driven on the critical path by
CompiledPretrainingStep._eager_step (not eager) and the loss cotangent
bridge IS wired and ready; the run reached the backward route and timed
out at segment region ..._chain_10_11 / op=m2rnn_bwd. After this fix the
m2rnn_bwd segment is time-chunked (8 launches, ~3s, watchdog-clear) and the
standalone direct-chain executor test runs both recurrent backward scans
(m2rnn_bwd seg 3 + mamba3_mimo_bwd seg 5) to completion.

Add _TIME_CHUNKED_RECURRENT_BACKWARD_OPS = {m2rnn_bwd, mamba3_mimo_bwd}
and key the planner launcher_chunks switch on membership instead of the
hard-coded mamba3-only check. Per-row-INDEPENDENT backward ops
(attention_qkv_projection_bwd / sparse_mla_fp8_apply_bwd / *_rmsnorm_bwd)
stay grid_chunks. Update the standalone executor test to assert both
recurrent backward segments are time-chunked.

NOTE (separate pre-existing defect, NOT fixed here, now UNMASKED): at FULL
local_gb10_quarter size the mamba3_mimo_bwd launcher_chunks kernel crashes
the Metal compiler service (XPC_ERROR_CONNECTION_INTERRUPTED during
newComputePipelineState) in a cold process too; the m2rnn watchdog timeout
previously hid it. The guarded runtime RAISES loudly (RULE #1), so the m04
critical path still does not complete end-to-end at full scale.

Author: apstenku123

cppmega_mlx/runtime/path_c_fusion_schedules.py

@@ -137,6 +137,30 @@
         "mamba3_mimo_bwd",
     }
 )
+# Recurrent REVERSE-TIME-SCAN backward ops. Each is a single
+# ``for time_rev in T.serial(0, S)`` over the WHOLE sequence inside ONE
+# threadgroup. Encoded under the default grid_chunks dispatch as ONE Metal
+# command buffer that spans all S time steps, the scan exceeds the macOS GPU
+# watchdog window (kIOGPUCommandBufferCallbackErrorTimeout). These ops MUST be
+# switched to launcher_chunks so the reverse scan is driven as K command buffers
+# over TIME (one ``artifact()`` per row-window, reverse adjoint state carried
+# across launches through caller-owned buffers). The other row-phased backward
+# ops in _EXACT_ROW_PHASED_BACKWARD_OPS (attention_qkv_projection_bwd /
+# sparse_mla_fp8_apply_bwd) are per-row-INDEPENDENT (no carried reverse-time
+# state) so they fit one command buffer and stay grid_chunks.
+#
+# NOTE: the original intra-segment time-chunking (commit ac412fb) only covered
+# mamba3_mimo_bwd, but in the direct-chain reverse order the m2rnn_bwd segment
+# runs BEFORE mamba3 and is an equally long reverse-time scan -> it tripped the
+# watchdog first (verified: timeout at segment region ..._chain_10_11 /
+# op=m2rnn_bwd, before mamba3 was ever reached). Both recurrent backward ops
+# need the launcher split.
+_TIME_CHUNKED_RECURRENT_BACKWARD_OPS = frozenset(
+    {
+        "m2rnn_bwd",
+        "mamba3_mimo_bwd",
+    }
+)
 DESCRIPTOR_ROW_PHASED_BWD_SCRATCH_ABI_BUFFERS = frozenset(
     {
         "attention_hidden",
@@ -14232,28 +14256,37 @@ def plan_path_c_direct_fusion_chain_for_region(
                 candidate_region,
                 DESCRIPTOR_PHYSICAL_ABI_POLICY_DIRECT,
             )
-            # Intra-segment TIME-CHUNKING for the mamba3 mimo BACKWARD segment.
+            # Intra-segment TIME-CHUNKING for the recurrent reverse-time-scan
+            # BACKWARD segments (m2rnn_bwd AND mamba3_mimo_bwd; see
+            # _TIME_CHUNKED_RECURRENT_BACKWARD_OPS).
             #
-            # The mamba3 reverse-time recurrent scan is a single
+            # Each recurrent reverse scan is a single
             # `for time_rev in T.serial(0, S)` over the whole sequence inside ONE
             # threadgroup. Under the default grid_chunks dispatch that is ONE Metal
-            # command buffer spanning all S time steps -> ~7s GPU time -> macOS GPU
-            # watchdog kill (kIOGPUCommandBufferCallbackErrorTimeout).
+            # command buffer spanning all S time steps -> multi-second GPU time ->
+            # macOS GPU watchdog kill (kIOGPUCommandBufferCallbackErrorTimeout).
             #
-            # Switching JUST this segment to launcher_chunks makes its compiled
-            # kernel declare path_c_row_chunk_index / path_c_row_subchunk_index and
-            # emit the per-row reverse scan body
+            # Switching the segment to launcher_chunks makes its compiled kernel
+            # declare path_c_row_chunk_index / path_c_row_subchunk_index and emit
+            # the per-row reverse scan body
             # (`for time_rev in T.serial(row, row + 1)`), so the runtime can drive
-            # the reverse scan as K separate command buffers over TIME, carrying the
-            # reverse adjoint scan state (dh<->h0_grad, angle_grad<->
-            # mamba3_angle_grad_state) across launches via caller-owned buffers.
-            # Every other segment (forward + non-mamba3-bwd) keeps grid_chunks --
-            # only the long reverse scan needs the watchdog-safe launcher split.
+            # the reverse scan as K separate command buffers over TIME, carrying
+            # the reverse adjoint scan state across launches via caller-owned
+            # buffers. Every other segment (forward + per-row-INDEPENDENT backward
+            # ops like attention_qkv_projection_bwd / sparse_mla_fp8_apply_bwd /
+            # residual_rmsnorm_bwd) keeps grid_chunks -- only the long reverse
+            # scans need the watchdog-safe launcher split.
+            #
+            # The original time-chunking (commit ac412fb) covered ONLY
+            # mamba3_mimo_bwd; m2rnn_bwd runs FIRST in the reverse chain and tripped
+            # the watchdog before mamba3 was reached (verified timeout at region
+            # ..._chain_10_11 / op=m2rnn_bwd), so both recurrent ops are included.
             if (
                 execution_phase == DESCRIPTOR_EXECUTION_STAGE_BACKWARD
                 and direct_target.max_rows_per_launch is not None
                 and any(
-                    _path_c_descriptor_stage_node_op_name(node) == "mamba3_mimo_bwd"
+                    _path_c_descriptor_stage_node_op_name(node)
+                    in _TIME_CHUNKED_RECURRENT_BACKWARD_OPS
                     for node in candidate_region.nodes
                 )
             ):

tests/test_m04_train_step.py

@@ -3624,30 +3624,37 @@ def test_path_c_direct_chain_runtime_executor_runs_native_segments(
         "ok",
         "ok",
     ]
-    # The mamba3 mimo BACKWARD segment is TIME-CHUNKED (launcher_chunks): its
-    # kernel additionally declares the path_c_row_chunk_index /
+    # The recurrent reverse-time-scan BACKWARD segments are TIME-CHUNKED
+    # (launcher_chunks): both the m2rnn_bwd (segment 3) and mamba3_mimo_bwd
+    # (segment 5) reverse scans declare the path_c_row_chunk_index /
     # path_c_row_subchunk_index / path_c_backward_stage_index scalar params so the
-    # runtime can split the reverse-time scan into watchdog-safe per-launch
-    # command buffers. That segment therefore carries 3 extra scalar args beyond
-    # the backward gate; every other segment keeps grid_chunks dispatch.
+    # runtime can split each reverse-time scan into watchdog-safe per-launch
+    # command buffers. Those segments therefore carry 3 extra scalar args beyond
+    # the backward gate; every other segment (forward + per-row-INDEPENDENT
+    # backward ops) keeps grid_chunks dispatch.
+    #
+    # m2rnn_bwd was added to the time-chunked set because in the reverse chain it
+    # runs BEFORE mamba3 and is an equally long reverse-time scan; under
+    # grid_chunks it tripped the macOS GPU watchdog before mamba3 was reached.
     time_chunked_segment_indices = {
         segment["index"]
         for segment in payload["segments"]
         if segment.get("time_chunk_launch_count")
     }
-    assert time_chunked_segment_indices == {5}, time_chunked_segment_indices
-    mamba3_bwd_launch_count = next(
-        segment["time_chunk_launch_count"]
-        for segment in payload["segments"]
-        if segment["index"] == 5
-    )
-    assert mamba3_bwd_launch_count > 1
+    assert time_chunked_segment_indices == {3, 5}, time_chunked_segment_indices
+    for time_chunked_index in sorted(time_chunked_segment_indices):
+        recurrent_bwd_launch_count = next(
+            segment["time_chunk_launch_count"]
+            for segment in payload["segments"]
+            if segment["index"] == time_chunked_index
+        )
+        assert recurrent_bwd_launch_count > 1
     assert [segment["kernel_arg_count"] for segment in payload["segments"]] == [
         len(rb_segment["required_logical_buffers"])
         # backward segments additionally pass the path_c_run_backward gate scalar
         + (1 if payload_segment["execution_phase"] == "backward" else 0)
-        # the time-chunked mamba3 backward also binds the launcher-chunk index
-        # scalars (path_c_row_chunk_index + path_c_row_subchunk_index +
+        # the time-chunked recurrent backward segments also bind the launcher-chunk
+        # index scalars (path_c_row_chunk_index + path_c_row_subchunk_index +
         # path_c_backward_stage_index) so the runtime can select each launch.
         + (3 if payload_segment["index"] in time_chunked_segment_indices else 0)
         for payload_segment, rb_segment in zip(