fix(path-c): clear Metal mamba3_mimo_bwd watchdog gate via per-op time-chunk window

The mamba3 reverse-time-scan backward segment tripped the macOS GPU watchdog
(kIOGPUCommandBufferCallbackErrorTimeout) on its FIRST time-chunk launch even
when time-chunked: the shared launcher window of 8 reverse time-steps/command
buffer is too big for mamba3, whose reverse-scan state (1.75MB/step x4) is pooled
to GLOBAL device memory (b330bdb) making each step expensive.

Add a PER-OP time-chunk window override so mamba3_mimo_bwd uses a smaller
rows_per_kernel_launch than the other launcher-chunked backward ops (which keep
the watchdog-safe shared default of 8). _target_with_max_rows_per_launch now
threads rows_per_kernel_launch into the generated schedule template (the window
lives only in the compiled PrimFunc attrs, so a non-default window always
re-generates the kernel -- no silent keep of the 8-step kernel).

Isolation sweep at full local_gb10_quarter scale (depth=13 hidden=3584
max_seq=4096) via scripts/_bwdgate_mamba3_window_sweep.py, first 4 launches:
  window  total launches  first launch  steady/launch  first-launch fits?
    8        512            (timeout)        --          NO (kIOGPU timeout)
    4       1024            3.63s          2.42s         yes (~1.4s margin)
    2       2048            2.52s          1.19s         yes (~2.5s margin)
    1       4096            1.65s          0.68s         yes (~3.3s margin)
Per-step cost is ~linear so total mamba3 backward wall is window-independent
(~40-46 min). 4 is the LARGEST watchdog-safe window; committed default is 2 for
~50% margin during the ~61-min end-to-end run. Metal-only; CUDA keeps the shared
default (no per-command-buffer watchdog).

Only the mamba3 segment changes (rows_per_kernel_launch 8->2 -> 2048 launches);
sparse_mla_fp8_apply_bwd / attention_qkv_projection_bwd / m2rnn_bwd keep window 8.
All 119 tests/test_path_c_fusion_ir.py pass.

Author: apstenku123

cppmega_mlx/runtime/path_c_fusion_schedules.py

@@ -278,6 +278,37 @@
 DESCRIPTOR_ROW_CHUNK_INDEX_PARAM = "path_c_row_chunk_index"
 DESCRIPTOR_ROW_SUBCHUNK_INDEX_PARAM = "path_c_row_subchunk_index"
 DESCRIPTOR_ROWS_PER_KERNEL_LAUNCH = 8
+# Per-op time-chunk window override for the mamba3 reverse-time-scan backward.
+#
+# Every launcher-chunked backward segment processes ``rows_per_kernel_launch``
+# reverse time-steps per Metal command buffer (one ``(chunk, subchunk)`` launch;
+# see _path_c_segment_time_chunk_launches). The shared default of 8 steps/launch
+# is watchdog-safe for the per-row-INDEPENDENT heavy ops
+# (attention_qkv_projection_bwd / sparse_mla_fp8_apply_bwd, ~0.42s/launch) but
+# NOT for mamba3_mimo_bwd: its reverse-scan state (1.75MB/step x 4) was pooled to
+# GLOBAL device memory (commit b330bdb, forced by Metal's 32KiB threadgroup cap),
+# so each reverse step is far slower and an 8-step command buffer trips the macOS
+# GPU watchdog (kIOGPUCommandBufferCallbackErrorTimeout) on its FIRST launch at
+# full local_gb10_quarter scale (depth=13 hidden=3584 max_seq=4096). mamba3 gets
+# its OWN smaller window here.
+#
+# Measured isolation sweep (scripts/_bwdgate_mamba3_window_sweep.py, M4 Max, full
+# local_gb10_quarter, first 4 launches; first launch carries the checkpoint-replay
+# setup, ~5s GPU watchdog window):
+#   window  total launches   first launch   steady/launch   first-launch fits?
+#   ------  --------------   ------------   -------------   ------------------
+#     8         512            (timeout)         --           NO  (kIOGPU timeout)
+#     4        1024             3.63s           2.42s         yes (tight, ~1.4s margin)
+#     2        2048             2.52s           1.19s         yes (~2.5s margin)
+#     1        4096             1.65s           0.68s         yes (~3.3s margin)
+# Per-step cost is ~linear (~0.6s/time-step) so the TOTAL mamba3 backward wall is
+# roughly window-independent (~40-46 min); reducing the window only trades a fixed
+# total cost for more launches with more watchdog margin. 4 is the LARGEST
+# watchdog-safe window; 2 is chosen as the committed default for a comfortable
+# ~50% margin during the ~61-min end-to-end run (a window-4 launch at 3.63s leaves
+# little headroom if a transient GPU stall lands on it). Metal-only: CUDA has no
+# per-command-buffer watchdog and keeps the shared default.
+MAMBA3_BWD_ROWS_PER_KERNEL_LAUNCH = 2
 DESCRIPTOR_BACKWARD_GATE_PARAM = "path_c_run_backward"
 DESCRIPTOR_BACKWARD_STAGE_INDEX_PARAM = "path_c_backward_stage_index"
 DESCRIPTOR_EXECUTION_STAGE_ALL = "all"
@@ -707,6 +738,27 @@ def _path_c_descriptor_stage_rows_per_kernel_launch_for_node(node: Any) -> int:
     return DESCRIPTOR_ROWS_PER_KERNEL_LAUNCH
 
 
+def _mamba3_bwd_rows_per_kernel_launch_for_nodes(nodes: Iterable[Any]) -> int:
+    """Per-launch time-chunk window for a launcher-chunked backward segment.
+
+    Returns :data:`MAMBA3_BWD_ROWS_PER_KERNEL_LAUNCH` (the smaller, watchdog-safe
+    mamba3 window) when the segment contains ``mamba3_mimo_bwd`` -- whose pooled
+    global reverse-scan state makes each reverse step expensive enough that the
+    shared 8-step window trips the macOS GPU watchdog. Every other
+    launcher-chunked backward op (the per-row-INDEPENDENT heavy ops and m2rnn_bwd)
+    keeps the shared :data:`DESCRIPTOR_ROWS_PER_KERNEL_LAUNCH` window, which is
+    already watchdog-safe for them. With the backward op cap = 1 each backward op
+    is in its OWN segment, so a segment never mixes mamba3 with another op.
+    """
+
+    if any(
+        _path_c_descriptor_stage_node_op_name(node) == "mamba3_mimo_bwd"
+        for node in nodes
+    ):
+        return MAMBA3_BWD_ROWS_PER_KERNEL_LAUNCH
+    return DESCRIPTOR_ROWS_PER_KERNEL_LAUNCH
+
+
 def _path_c_descriptor_stage_append(
     groups: list[PathCDescriptorScheduleStageGroup],
     *,
@@ -14817,11 +14869,23 @@ def plan_path_c_direct_fusion_chain_for_region(
                     for node in candidate_region.nodes
                 )
             ):
+                # Per-op time-chunk window. mamba3_mimo_bwd carries multi-MB
+                # global reverse-scan state per step (b330bdb), so an 8-step
+                # command buffer trips the watchdog on its first launch -- it
+                # needs a SMALLER window than the per-row-independent heavy ops.
+                # _mamba3_bwd_rows_per_kernel_launch_for_nodes returns the mamba3
+                # override only when the segment contains mamba3_mimo_bwd; every
+                # other launcher-chunked backward op keeps the shared default.
                 direct_target = _target_with_max_rows_per_launch(
                     direct_target,
                     candidate_region,
                     DESCRIPTOR_DEFAULT_MAX_ROWS_PER_LAUNCH,
                     DESCRIPTOR_ROW_DISPATCH_LAUNCHER_CHUNKS,
+                    rows_per_kernel_launch=(
+                        _mamba3_bwd_rows_per_kernel_launch_for_nodes(
+                            candidate_region.nodes
+                        )
+                    ),
                 )
             # ROW-WINDOWING for the per-row-INDEPENDENT heavy FORWARD ops
             # (attention_qkv_projection / sparse_mla_fp8_apply; see
@@ -15092,12 +15156,25 @@ def _target_with_max_rows_per_launch(
     region: PathCFusionRegion,
     max_rows_per_launch: int | None,
     row_dispatch_mode: str,
+    rows_per_kernel_launch: int = DESCRIPTOR_ROWS_PER_KERNEL_LAUNCH,
 ) -> PathCFusionScheduleTarget:
     validated_rows = _validated_max_rows_per_launch(max_rows_per_launch)
     validated_mode = _validated_row_dispatch_mode(row_dispatch_mode)
+    validated_rows_per_kernel_launch = _validated_rows_per_kernel_launch(
+        rows_per_kernel_launch
+    )
+    # The dispatch mode + max_rows_per_launch are carried on the target struct,
+    # but the per-launch time-chunk window (rows_per_kernel_launch) lives ONLY in
+    # the generated schedule_template / compiled PrimFunc attrs, so it is never
+    # recorded on the PathCFusionScheduleTarget. A non-default window must
+    # therefore re-generate the template even when mode + max_rows already match;
+    # the early-out is valid only for the default window (where re-generation
+    # would be a no-op). RULE #1: an explicit non-default window always lowers --
+    # never silently keeps the default 8-step kernel.
     if (
         target.max_rows_per_launch == validated_rows
         and target.row_dispatch_mode == validated_mode
+        and validated_rows_per_kernel_launch == DESCRIPTOR_ROWS_PER_KERNEL_LAUNCH
     ):
         return target
     shape_env = _shape_env_for_region(region)
@@ -15114,6 +15191,7 @@ def _target_with_max_rows_per_launch(
             physical_abi_policy=target.physical_abi_policy,
             max_rows_per_launch=validated_rows,
             row_dispatch_mode=validated_mode,
+            rows_per_kernel_launch=validated_rows_per_kernel_launch,
         ),
         max_rows_per_launch=validated_rows,
         row_dispatch_mode=validated_mode,

scripts/_bwdgate_mamba3_window_sweep.py

@@ -0,0 +1,200 @@
+#!/usr/bin/env python3
+"""Isolation sweep for the mamba3_mimo_bwd time-chunk window at FULL scale.
+
+Drives ONLY the mamba3_mimo_bwd backward segment of the full local_gb10_quarter
+Path C direct-chain (depth=13, hidden=3584, max_seq=4096) -- NO 61-min forward.
+For a given ``--rows-per-kernel-launch`` window it:
+
+  1. monkeypatches MAMBA3_BWD_ROWS_PER_KERNEL_LAUNCH so the mamba3 segment is
+     compiled with that per-launch time-chunk window,
+  2. plans + compiles the chain, extracts the mamba3 segment,
+  3. allocates caller-owned zero buffers,
+  4. drives the segment's time-chunk launches ONE AT A TIME, timing EACH launch
+     (eval + synchronize == one committed Metal command buffer), capturing the
+     GPU watchdog timeout per-launch so we attribute the failure to the FIRST
+     launch.
+
+The FIRST launch (chunk 0, subchunk 0) is the watchdog gate: it processes the
+LAST ``window`` reverse time-steps AND runs the checkpoint-replay for them.
+
+By default only the first ``--max-launches`` launches are timed (enough to see
+whether the first launch fits the watchdog and the steady per-launch cost);
+pass ``--all-launches`` to run the ENTIRE mamba3 backward (4096/window launches)
+and report the total mamba3-backward wall time.
+
+RULE #1: no silent fallback. A watchdog timeout on the first launch is RAISED
+and reported as that window being infeasible.
+"""
+from __future__ import annotations
+
+import argparse
+import dataclasses
+import json
+import sys
+import time
+import traceback
+from pathlib import Path
+from typing import Any
+
+ROOT = Path(__file__).resolve().parents[1]
+for p in (str(ROOT), str(ROOT / "scripts")):
+    if p not in sys.path:
+        sys.path.insert(0, p)
+
+import mlx.core as mx  # noqa: E402
+import m04_train_step as m  # noqa: E402
+import cppmega_mlx.runtime.path_c_fusion_schedules as sched  # noqa: E402
+
+
+def _find_mamba3_segment(chain: Any) -> Any:
+    for seg in chain.segments:
+        if any(n.op_name == "mamba3_mimo_bwd" for n in seg.region.nodes):
+            return seg
+    raise RuntimeError("no mamba3_mimo_bwd segment found in direct-chain plan")
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--rows-per-kernel-launch", type=int, required=True,
+                        help="time-steps per launch for the mamba3 bwd window")
+    parser.add_argument("--max-launches", type=int, default=4,
+                        help="number of leading launches to drive+time (gate probe)")
+    parser.add_argument("--all-launches", action="store_true",
+                        help="drive ALL launches (full mamba3 backward); reports total")
+    args = parser.parse_args()
+
+    window = int(args.rows_per_kernel_launch)
+    if window <= 0:
+        raise ValueError("--rows-per-kernel-launch must be positive")
+
+    # Force the mamba3 per-op time-chunk window for THIS run.
+    sched.MAMBA3_BWD_ROWS_PER_KERNEL_LAUNCH = window
+
+    print(f"[sweep] device={mx.default_device()} target={m._path_c_default_target()}",
+          flush=True)
+    print(f"[sweep] mamba3 rows_per_kernel_launch (window) = {window}", flush=True)
+
+    profile, route_symbols, regions = m._local_gb10_path_c_model_regions()
+    print(f"[sweep] profile={profile.name} depth={profile.depth} "
+          f"hidden={profile.hidden_size} max_seq={profile.max_seq_length}", flush=True)
+    sel = m._select_path_c_model_route_region(regions)
+    scheduled = m.plan_path_c_fusion_schedule_for_region(sel, include_backward=True)
+    chain = m.plan_path_c_direct_fusion_chain_for_region(
+        scheduled.region, include_backward=True,
+    )
+    mamba_seg = _find_mamba3_segment(chain)
+    tgt = mamba_seg.schedule_target
+    print(f"[sweep] mamba3 seg index={mamba_seg.index} region={mamba_seg.region.name} "
+          f"row_dispatch={getattr(tgt, 'row_dispatch_mode', None)} "
+          f"max_rows_per_launch={getattr(tgt, 'max_rows_per_launch', None)}", flush=True)
+
+    # Compile ALL segments (the runtime route needs the full artifact set), then
+    # drive ONLY the mamba3 sub-chain.
+    t_c = time.perf_counter()
+    artifacts = m.compile_path_c_direct_fusion_chain_artifacts(chain)
+    print(f"[sweep] compiled {len(artifacts)} segment shaders in "
+          f"{time.perf_counter()-t_c:.1f}s", flush=True)
+
+    # Inspect the compiled mamba3 prim_func to report the launch math.
+    prim_func = None
+    for art in artifacts:
+        pf = getattr(art, "prim_func", None) or getattr(art, "_prim_func", None)
+        # artifacts is a dict in the per-segment probe; handle both.
+    # The route runner reads launch math off the prim_func; surface it via the
+    # same helper the runtime uses.
+    sub_chain = dataclasses.replace(chain, segments=(mamba_seg,))
+
+    # Hook per-launch timing by wrapping mx.synchronize (each time-chunk launch
+    # commits via eval()+synchronize() -- the synchronize is the command-buffer
+    # boundary). We time the interval between consecutive synchronize() returns.
+    launch_times: list[float] = []
+    state: dict[str, Any] = {"last": None}
+    real_sync = mx.synchronize
+
+    def timed_sync(*a: Any, **k: Any) -> Any:
+        r = real_sync(*a, **k)
+        now = time.perf_counter()
+        if state["last"] is not None:
+            dt = now - state["last"]
+            launch_times.append(dt)
+            print(f"[sweep]   launch[{len(launch_times)-1}] committed in "
+                  f"{dt:.3f}s", flush=True)
+        state["last"] = now
+        return r
+
+    mx.synchronize = timed_sync  # type: ignore[assignment]
+
+    # Limit the number of launches when not running the full backward, by
+    # truncating the launch list the runtime iterates. We do that by patching
+    # _path_c_segment_time_chunk_launches in m04 to slice to max_launches.
+    real_launches_fn = m._path_c_segment_time_chunk_launches
+    limit = None if args.all_launches else int(args.max_launches)
+
+    def limited_launches(pf: Any) -> tuple:
+        full = real_launches_fn(pf)
+        if not full:
+            return full
+        print(f"[sweep] mamba3 total time-chunk launches (full) = {len(full)} "
+              f"(window={window}, S/window)", flush=True)
+        if limit is None:
+            return full
+        return full[:limit]
+
+    m._path_c_segment_time_chunk_launches = limited_launches  # type: ignore[assignment]
+
+    specs = m._path_c_direct_chain_required_logical_buffer_specs(sub_chain)
+    buffers: dict[str, Any] = {}
+    for name, spec in specs.items():
+        dtype = getattr(mx, str(spec["dtype"]))
+        buffers[name] = mx.zeros(tuple(int(d) for d in spec["shape"]), dtype=dtype)
+    mx.eval(*buffers.values())
+    print(f"[sweep] allocated {len(buffers)} caller-owned buffers; driving "
+          f"{'ALL' if args.all_launches else limit} mamba3 launches", flush=True)
+
+    status = "ok"
+    err = None
+    t0 = time.perf_counter()
+    # Seed the per-launch clock so launch[0]'s window (run-start -> first
+    # synchronize) is measured. The route runner's pre-launch buffer eval is
+    # cheap (zero buffers already materialised); the first synchronize boundary
+    # therefore measures launch[0]'s GPU command buffer.
+    state["last"] = t0
+    try:
+        payload = m.run_path_c_direct_fusion_chain_route(
+            chain=sub_chain,
+            logical_buffers=buffers,
+            artifacts=artifacts,
+        )
+    except BaseException as exc:  # noqa: BLE001
+        status = "FAIL"
+        err = f"{type(exc).__name__}: {exc}"
+        traceback.print_exc()
+        payload = None
+    finally:
+        mx.synchronize = real_sync  # type: ignore[assignment]
+        m._path_c_segment_time_chunk_launches = real_launches_fn  # type: ignore[assignment]
+    total = time.perf_counter() - t0
+
+    first_launch = launch_times[0] if launch_times else None
+    steady = (
+        sum(launch_times[1:]) / max(1, len(launch_times) - 1)
+        if len(launch_times) > 1 else None
+    )
+    out = {
+        "window_rows_per_kernel_launch": window,
+        "status": status,
+        "error": err,
+        "driven_launches": (None if args.all_launches else limit),
+        "all_launches": bool(args.all_launches),
+        "timed_launch_count": len(launch_times),
+        "first_launch_s": round(first_launch, 4) if first_launch is not None else None,
+        "steady_per_launch_s": round(steady, 4) if steady is not None else None,
+        "wall_total_s": round(total, 3),
+        "first_launch_fits_watchdog": (status == "ok"),
+    }
+    print("\n[sweep] RESULT:\n" + json.dumps(out, indent=2), flush=True)
+    return 0 if status == "ok" else 7
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())