feat(path-c): intra-segment TIME-CHUNKING for mamba3 mimo backward (Metal watchdog fix)

The mamba3 mimo BACKWARD direct-chain segment is a reverse-time recurrent scan
over the whole sequence (S serial time steps) inside ONE threadgroup. Under the
default grid_chunks dispatch the entire scan is encoded into ONE Metal command
buffer; once its GPU time exceeds the macOS watchdog window it is killed with
kIOGPUCommandBufferCallbackErrorTimeout (verified: monolithic seg dies at 6.35s
for S=16384; time-chunked completes clean).

Two-edit fix (Python only, no tilelang C++ rebuild):

1. plan_path_c_direct_fusion_chain_for_region: switch ONLY the mamba3_mimo_bwd
   backward segment's schedule_target to row_dispatch_mode=launcher_chunks
   (max_rows_per_launch=64). Its compiled kernel then declares
   path_c_row_chunk_index / path_c_row_subchunk_index and emits the per-row
   reverse scan body (for time_rev in serial(row, row+1)). Every other segment
   (all forward + non-mamba3 backward) keeps grid_chunks unchanged.

   _kernel_parameter_count_for_target now counts only params that bind a device
   buffer (present in buffer_map), not scalar params. The portable limit (31) is
   Metal's buffer-argument cap; the launcher index scalars are passed by value
   and consume no buffer slot. Without this the launcher-chunked mamba3 backward
   was spuriously blocked at 34 'buffers' when it binds only 30.

2. run_path_c_direct_fusion_chain_route: for a launcher-chunked segment on Metal,
   loop over (chunk_index, subchunk_index) in [0,S) (mirroring the eager replay
   driver order), write the two index scalars into a per-launch copy of the args
   at their recorded segment-dependent positions, call the artifact per launch,
   and commit each launch as its own command buffer (eval+synchronize). The
   reverse adjoint scan state (dh<->h0_grad, angle_grad<->mamba3_angle_grad_state)
   carries across launches via caller-owned buffers (relaxed atomic_add already
   accumulates correctly); path_c_first_row_launch zeroes owner grads exactly
   once on the first launch. CUDA branch unchanged; non-launcher-chunked segments
   keep the single-call path. No fallback: missing index scalar slots or a commit
   failure RAISE.

Verification (Metal, S=512 bench shape and beyond):
 - native executor test runs all 7 segments incl. the time-chunked seg#5: PASS.
 - gradient parity time-chunked vs monolithic: worst mamba3 grad diff 5.59e-5,
   within the inherent fp8/bf16 atomic run-to-run noise floor (4.71e-5); no NaN.
 - watchdog: monolithic mamba3 backward raises kIOGPUCommandBufferCallbackError
   Timeout at 6.35s (S=16384); time-chunked (2048 launches, 3.84ms each)
   completes in 7.86s with NO timeout.

No new test regressions (the 5 pre-existing failures in test_m04_train_step.py /
test_path_c_fusion_ir.py fail identically on the pristine baseline).

Author: apstenku123

cppmega_mlx/runtime/path_c_fusion_schedules.py

@@ -14232,6 +14232,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.
+            #
+            # The mamba3 reverse-time recurrent 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).
+            #
+            # 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
+            # (`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.
+            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"
+                    for node in candidate_region.nodes
+                )
+            ):
+                direct_target = _target_with_max_rows_per_launch(
+                    direct_target,
+                    candidate_region,
+                    DESCRIPTOR_DEFAULT_MAX_ROWS_PER_LAUNCH,
+                    DESCRIPTOR_ROW_DISPATCH_LAUNCHER_CHUNKS,
+                )
             try:
                 parameter_count = _kernel_parameter_count_for_target(
                     candidate_region,
@@ -14483,7 +14514,19 @@ def _kernel_parameter_count_for_target(
     region: PathCFusionRegion,
     target: PathCFusionScheduleTarget,
 ) -> int:
-    return len(tuple(target.schedule_template(region).params))
+    # The portable limit (DESCRIPTOR_PORTABLE_KERNEL_BUFFER_LIMIT) is Metal's
+    # buffer-argument binding cap, so count only params that actually bind a
+    # device buffer (present in the prim_func ``buffer_map``). Scalar params --
+    # the backward gate plus the launcher time-chunk index params
+    # (path_c_row_chunk_index / path_c_row_subchunk_index /
+    # path_c_backward_stage_index) -- are passed by value and consume NO buffer
+    # slot, so they must not count against the buffer limit. (Counting them
+    # spuriously blocked the launcher-chunked mamba3 backward at 34 "buffers"
+    # when it really binds only 30.)
+    prim_func = target.schedule_template(region)
+    params = tuple(getattr(prim_func, "params", ()))
+    buffer_map = getattr(prim_func, "buffer_map", {}) or {}
+    return sum(1 for param in params if buffer_map.get(param) is not None)
 
 
 def _attested_schedule_template_for_target(

scripts/m04_train_step.py

@@ -6497,6 +6497,65 @@ def _path_c_call_cuda_artifact_with_mlx_bridge(
     return result
 
 
+def _path_c_segment_time_chunk_launches(prim_func: Any) -> tuple[tuple[int, int], ...]:
+    """Return ``(chunk_index, subchunk_index)`` launch pairs for a segment.
+
+    A direct-chain segment whose compiled prim_func declares
+    ``_cppmega_path_c_max_rows_per_launch`` AND
+    ``_cppmega_path_c_row_dispatch_mode == 'launcher_chunks'`` is time-chunked:
+    its reverse-time recurrent scan body emits one time step per ``row`` and
+    the runtime selects which rows a launch processes via the
+    ``path_c_row_chunk_index`` / ``path_c_row_subchunk_index`` scalar params.
+    The kernel maps ``(chunk_index, subchunk_index)`` to a row window and walks
+    only those rows, carrying the reverse scan state (dh<->h0_grad,
+    angle_grad<->mamba3_angle_grad_state) across launches through caller-owned
+    buffers; ``path_c_first_row_launch`` (1 only at chunk 0 / subchunk 0) zeroes
+    the owner grads exactly once.
+
+    Iteration order MIRRORS the eager replay driver
+    (cppmega_mlx.training.path_c_fused_replay.replay_launch_scalar_params):
+    ``for chunk in range(row_chunk_count): for subchunk in
+    range(row_subchunk_count)``. Because the reverse-scan body computes
+    ``time_idx = (S - 1) - row``, the first launch (chunk 0, subchunk 0, row 0)
+    correctly starts the reverse scan at the LAST time step (S-1), so forward
+    chunk/subchunk order walks time S-1 -> 0 exactly as the monolithic scan did.
+
+    Returns an empty tuple when the segment is NOT launcher-chunked (grid_chunks
+    or non-row-phased) -- the caller then keeps the single-launch path, which is
+    the correct clear path for those segments (NOT a fallback).
+    """
+
+    max_rows_per_launch = getattr(
+        prim_func, "_cppmega_path_c_max_rows_per_launch", None
+    )
+    row_dispatch_mode = getattr(prim_func, "_cppmega_path_c_row_dispatch_mode", None)
+    row_chunk_index_param = getattr(
+        prim_func, "_cppmega_path_c_row_chunk_index_param", None
+    )
+    if (
+        max_rows_per_launch is None
+        or str(row_dispatch_mode) != "launcher_chunks"
+        or not row_chunk_index_param
+    ):
+        return ()
+    row_chunk_count = getattr(prim_func, "_cppmega_path_c_row_chunk_count", None)
+    row_subchunk_count = getattr(
+        prim_func, "_cppmega_path_c_row_subchunk_count", None
+    )
+    if row_chunk_count is None:
+        raise ValueError(
+            "launcher-chunked direct-chain segment is missing "
+            "_cppmega_path_c_row_chunk_count"
+        )
+    chunk_count = max(1, int(row_chunk_count))
+    subchunk_count = max(1, int(row_subchunk_count or 1))
+    return tuple(
+        (chunk_index, subchunk_index)
+        for chunk_index in range(chunk_count)
+        for subchunk_index in range(subchunk_count)
+    )
+
+
 def run_path_c_direct_fusion_chain_route(
     *,
     chain: Any,
@@ -6626,6 +6685,11 @@ def run_path_c_direct_fusion_chain_route(
         # the kernel's in-place results back into the owning bank buffers
         # (Metal is zero-copy so it does not need this).
         kernel_arg_buffer_names: dict[int, tuple[str, tuple[int, ...] | None]] = {}
+        # Record the positional index of every scalar param the time-chunk
+        # driver overrides per launch. The launcher-chunk index params sit at a
+        # SEGMENT-DEPENDENT position (they are interleaved with buffer params in
+        # kernel_param_order), so recompute them per segment -- never hardcode.
+        scalar_param_arg_positions: dict[str, int] = {}
         for name in kernel_param_order:
             if name in buffers:
                 kernel_arg_buffer_names[len(kernel_call_args)] = (
@@ -6642,6 +6706,7 @@ def run_path_c_direct_fusion_chain_route(
             elif name == gate_param:
                 kernel_call_args.append(run_backward_value)
             elif name in PATH_C_SCALAR_KERNEL_PARAM_DEFAULTS:
+                scalar_param_arg_positions[str(name)] = len(kernel_call_args)
                 kernel_call_args.append(PATH_C_SCALAR_KERNEL_PARAM_DEFAULTS[name])
             else:
                 raise ValueError(
@@ -6669,6 +6734,95 @@ def run_path_c_direct_fusion_chain_route(
                 mx_module=mx_module,
             )
         else:
+            # --- Intra-segment TIME-CHUNKING for the launcher-chunked segment ----
+            # The mamba3 mimo BACKWARD segment is a reverse-time recurrent scan
+            # over the WHOLE sequence (S time steps). Encoded as ONE Metal command
+            # buffer it runs ~7s of GPU time -> macOS GPU watchdog kill
+            # (kIOGPUCommandBufferCallbackErrorTimeout). Its compiled kernel
+            # declares path_c_row_chunk_index / path_c_row_subchunk_index and
+            # emits a per-row reverse scan body, so the runtime splits the scan
+            # into K command buffers over TIME: one artifact() call per
+            # (chunk, subchunk), each processing <= rows_per_kernel_launch time
+            # steps, with eval()+synchronize() committing each as its own command
+            # buffer. The reverse adjoint scan state (dh<->h0_grad,
+            # angle_grad<->mamba3_angle_grad_state) carries across launches via the
+            # caller-owned buffers (already accumulated with relaxed atomic_add),
+            # and path_c_first_row_launch zeroes the owner grads exactly once on
+            # the very first (chunk 0, subchunk 0) launch.
+            #
+            # Segments that are NOT launcher-chunked (every forward segment plus
+            # the non-mamba3 backward segments stay grid_chunks) take the single
+            # artifact() call below -- that is the correct clear path for them, not
+            # a fallback.
+            time_chunk_launches = _path_c_segment_time_chunk_launches(prim_func)
+            chunk_param = str(
+                getattr(prim_func, "_cppmega_path_c_row_chunk_index_param", "")
+                or ""
+            )
+            subchunk_param = str(
+                getattr(prim_func, "_cppmega_path_c_row_subchunk_index_param", "")
+                or ""
+            )
+            if time_chunk_launches:
+                # The launcher index params MUST be bound as kernel args for the
+                # split to select work; if the compiled kernel declares
+                # launcher-chunk dispatch but the args list lacks those scalar
+                # slots, the split would silently run the same window K times.
+                # RULE #1: fail loud rather than emit wrong grads.
+                if (
+                    chunk_param not in scalar_param_arg_positions
+                    or subchunk_param not in scalar_param_arg_positions
+                ):
+                    raise ValueError(
+                        f"direct-chain segment {segment.index} declares "
+                        "launcher-chunk dispatch but kernel args are missing the "
+                        f"{chunk_param!r}/{subchunk_param!r} scalar slots "
+                        f"(have {sorted(scalar_param_arg_positions)})"
+                    )
+                chunk_pos = scalar_param_arg_positions[chunk_param]
+                subchunk_pos = scalar_param_arg_positions[subchunk_param]
+                result = None
+                for chunk_index, subchunk_index in time_chunk_launches:
+                    launch_args = list(arrays)
+                    launch_args[chunk_pos] = int(chunk_index)
+                    launch_args[subchunk_pos] = int(subchunk_index)
+                    result = artifact(*launch_args)
+                    # Commit THIS launch as its own Metal command buffer so it
+                    # holds only <= rows_per_kernel_launch reverse-scan time steps
+                    # and stays under the watchdog window. eval()+synchronize() are
+                    # MLX's public scheduler entry points (commit on the scheduler
+                    # outer loop), so the cross-launch state handoff through the
+                    # caller-owned buffers is fully materialized before the next
+                    # launch reads it.
+                    mx_module.eval(*buffer_arrays)
+                    if _path_c_per_segment_command_buffer_commit_enabled():
+                        try:
+                            mx_module.synchronize()
+                        except Exception as exc:  # pragma: no cover
+                            raise RuntimeError(
+                                "direct-chain time-chunk command-buffer commit "
+                                f"failed at segment {int(segment.index)} "
+                                f"launch (chunk={chunk_index}, "
+                                f"subchunk={subchunk_index}) "
+                                f"(phase={execution_phase}, "
+                                f"region={segment.region.name}): "
+                                f"{type(exc).__name__}: {exc}"
+                            ) from exc
+                segment_results.append(
+                    {
+                        "index": int(segment.index),
+                        "status": "ok",
+                        "region_name": segment.region.name,
+                        "execution_phase": execution_phase,
+                        "schedule_id": target.schedule_id,
+                        "kernel_arg_count": len(arrays),
+                        "kernel_args": list(kernel_param_order),
+                        "elapsed_s": time.perf_counter() - segment_started,
+                        "result_type": type(result).__name__,
+                        "time_chunk_launch_count": len(time_chunk_launches),
+                    }
+                )
+                continue
             result = artifact(*arrays)
             # --- Per-stage Metal command-buffer commit boundary -----------------
             # The fused direct-chain encodes every segment's TVM-FFI kernel into

tests/test_m04_train_step.py

@@ -3624,10 +3624,32 @@ 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 /
+    # 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.
+    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 [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 +
+        # 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(
             payload["segments"],
             route["path_c_fusion"]["direct_chained_fusion"]["runtime_binding"][