DatasunriseOU
diff --git a/‎cppmega_mlx/runtime/path_c_fusion_schedules.py‎
Lines changed: 140 additions & 9 deletions b/‎cppmega_mlx/runtime/path_c_fusion_schedules.py‎
Lines changed: 140 additions & 9 deletions
diff --git a/‎scripts/_bwdgate_determinism_probe.py‎
Lines changed: 127 additions & 0 deletions b/‎scripts/_bwdgate_determinism_probe.py‎
Lines changed: 127 additions & 0 deletions
@@ -161,6 +161,39 @@
         "mamba3_mimo_bwd",
     }
 )
+# Per-row-INDEPENDENT heavy BACKWARD ops that exceed the macOS GPU watchdog
+# (~5-6s per command buffer) when run as ONE monolithic grid_chunks command
+# buffer at full local_gb10_quarter scale (depth=13 hidden=3584 max_seq=4096).
+#
+# Measured (full scale, each op ISOLATED in its own segment, grid_chunks, run in
+# a FRESH process so no prior GPU error pollutes the device):
+#   sparse_mla_fp8_apply_bwd      monolithic grid_chunks -> ~12s  -> WATCHDOG KILL
+#                                 (kIOGPUCommandBufferCallbackErrorTimeout)
+#   attention_qkv_projection_bwd  monolithic grid_chunks -> ~10s  -> WATCHDOG KILL
+#   residual_rmsnorm_bwd          monolithic grid_chunks -> ~0.08s -> safe
+# Both heavy ops time out the watchdog even ALONE (not from summing several ops
+# in one command buffer), so a fused-segment op cap alone does NOT clear them --
+# each needs its single op split across multiple command buffers over its
+# per-row axis. Unlike the recurrent reverse-time scans (m2rnn_bwd /
+# mamba3_mimo_bwd), these ops are per-row-INDEPENDENT: each output row depends
+# only on its own input row plus the shared weights, with NO carried reverse-time
+# state. So launcher_chunks row-windowing splits them into K command buffers over
+# rows with zero cross-launch state (see _row_phased_launcher_carry_buffers_for_nodes
+# which adds carry buffers ONLY for mamba3/m2rnn, never for these). The weight
+# gradients still accumulate correctly: path_c_first_row_launch zeroes the owner
+# grads exactly once on the first (chunk 0, subchunk 0) launch, then every launch
+# adds its rows' contribution.
+# Measured row-windowed (launcher_chunks, max_rows_per_launch=64 ->
+# rows_per_kernel_launch=8 -> 512 launches of 8 rows each):
+#   attention_qkv_projection_bwd  ~0.42s per launch -> watchdog-safe.
+# Metal-only (CUDA's compiler/scheduler has no such per-command-buffer watchdog,
+# so CUDA keeps the monolithic grid_chunks path).
+_ROW_CHUNKED_INDEPENDENT_BACKWARD_OPS = frozenset(
+    {
+        "sparse_mla_fp8_apply_bwd",
+        "attention_qkv_projection_bwd",
+    }
+)
 DESCRIPTOR_ROW_PHASED_BWD_SCRATCH_ABI_BUFFERS = frozenset(
     {
         "attention_hidden",
@@ -231,6 +264,31 @@
 # limit, so this cap is gated to Metal only (CUDA keeps greedy forward fusion).
 METAL_FORWARD_MAX_SEGMENT_NODES = 2
 
+# Metal-only BACKWARD fused-segment op-count cap.
+#
+# Sibling of METAL_FORWARD_MAX_SEGMENT_NODES, for the reverse-autograd backward
+# phase. The greedy direct-chain planner otherwise fuses backward ops up to the
+# portable buffer limit, producing the 3-op backward mega-kernel
+# ``sparse_mla_fp8_apply_bwd + attention_qkv_projection_bwd + residual_rmsnorm_bwd``
+# (region ``..._chain_7_10``). Measured at full local_gb10_quarter scale that
+# fused segment is ONE ~115 KB grid_chunks command buffer that runs ~10-25s of
+# GPU time and is killed by the macOS GPU watchdog
+# (kIOGPUCommandBufferCallbackErrorTimeout) -- the FIRST backward segment to die.
+# Capping backward segments at 1 op puts each backward op in its OWN segment /
+# command buffer, which (a) keeps each command buffer small and (b) lets the two
+# per-row-INDEPENDENT heavy ops (sparse_mla_fp8_apply_bwd,
+# attention_qkv_projection_bwd; see _ROW_CHUNKED_INDEPENDENT_BACKWARD_OPS) get
+# launcher_chunks row-windowing so each command buffer holds only a row window's
+# worth of work (~0.42s) -- well under the watchdog. The light backward ops
+# (residual_rmsnorm_bwd / entry_rmsnorm_bwd) run monolithically in ~0.08s.
+#
+# NOTE: an op cap alone does NOT clear sparse_mla_fp8_apply_bwd /
+# attention_qkv_projection_bwd -- each times out the watchdog even ISOLATED in a
+# 1-op segment, so isolation MUST be paired with row-windowing (the cap of 1
+# guarantees the isolation those ops need to be row-chunked). CUDA has no such
+# watchdog, so the cap is Metal-only (CUDA keeps greedy backward fusion).
+METAL_BACKWARD_MAX_SEGMENT_NODES = 1
+
 
 class _ResolveFromTarget:
     """Sentinel: resolve a planner default from the active lowering target.
@@ -246,6 +304,7 @@ def __repr__(self) -> str:  # pragma: no cover - debug aid
 
 
 _RESOLVE_FORWARD_CAP_FROM_TARGET = _ResolveFromTarget()
+_RESOLVE_BACKWARD_CAP_FROM_TARGET = _ResolveFromTarget()
 _GENERIC_MODEL_REAL_ABI_INPUT_SUFFIXES = (
     "residual_norm_weight",
     # Block A: the per-brick entry RMSNorm weight emitted by the model-
@@ -14230,6 +14289,9 @@ def plan_path_c_direct_fusion_chain_for_region(
     forward_max_segment_nodes: int | None | _ResolveFromTarget = (
         _RESOLVE_FORWARD_CAP_FROM_TARGET
     ),
+    backward_max_segment_nodes: int | None | _ResolveFromTarget = (
+        _RESOLVE_BACKWARD_CAP_FROM_TARGET
+    ),
     registry: PathCFusionScheduleRegistry | None = None,
 ) -> PathCFusionScheduleChainPlan:
     """Greedily split a Path C region into direct-buffer fused segments.
@@ -14249,11 +14311,22 @@ def plan_path_c_direct_fusion_chain_for_region(
     ``MTLCompilerService`` at ``newComputePipelineState`` with
     ``XPC_ERROR_CONNECTION_INTERRUPTED`` — splits into two smaller
     pipeline-safe kernels); on CUDA it is ``None`` (no cap, the CUDA compiler
-    has no such pipeline-state size limit, so CUDA fusion is unchanged). The
-    cap applies ONLY to forward segments; backward segments keep greedy fusion
-    up to the buffer limit (their reverse-time watchdog split is handled
-    separately via launcher_chunks). Pass an explicit value to override; pass
-    ``None`` to disable the cap entirely.
+    has no such pipeline-state size limit, so CUDA fusion is unchanged).
+
+    ``backward_max_segment_nodes`` is the sibling Metal-only cap for the
+    reverse-autograd BACKWARD phase. The default
+    (``_RESOLVE_BACKWARD_CAP_FROM_TARGET``) resolves to
+    :data:`METAL_BACKWARD_MAX_SEGMENT_NODES` = 1 on Metal and ``None`` on CUDA.
+    Capping backward segments at 1 op splits the 3-op backward mega-kernel
+    ``sparse_mla_fp8_apply_bwd + attention_qkv_projection_bwd +
+    residual_rmsnorm_bwd`` (region ``..._chain_7_10``, ~115 KB, ~10-25s GPU time
+    -> macOS GPU watchdog kill) into one segment per op, which both shrinks each
+    command buffer AND lets the per-row-INDEPENDENT heavy ops
+    (``sparse_mla_fp8_apply_bwd`` / ``attention_qkv_projection_bwd``;
+    :data:`_ROW_CHUNKED_INDEPENDENT_BACKWARD_OPS`) get launcher_chunks
+    row-windowing so each command buffer holds only a row window (~0.42s) -- well
+    under the watchdog. Pass an explicit value to override; pass ``None`` to
+    disable (CUDA keeps greedy backward fusion).
     """
 
     if not isinstance(region, PathCFusionRegion):
@@ -14279,6 +14352,23 @@ def plan_path_c_direct_fusion_chain_for_region(
         raise ValueError(
             "forward_max_segment_nodes must be positive when provided"
         )
+    if isinstance(backward_max_segment_nodes, _ResolveFromTarget):
+        # Resolve the backward fusion cap from the active lowering target: Metal
+        # caps at 1 op/segment (so the watchdog-killing 3-op backward segment
+        # ..._chain_7_10 splits per op and the heavy per-row-independent ops can
+        # be row-windowed); CUDA stays uncapped (no watchdog).
+        backward_max_segment_nodes = (
+            METAL_BACKWARD_MAX_SEGMENT_NODES
+            if _path_c_default_target() == "metal"
+            else None
+        )
+    if (
+        backward_max_segment_nodes is not None
+        and backward_max_segment_nodes <= 0
+    ):
+        raise ValueError(
+            "backward_max_segment_nodes must be positive when provided"
+        )
     working_region = (
         build_path_c_aot_autograd_region(region)
         if include_backward
@@ -14335,6 +14425,25 @@ def plan_path_c_direct_fusion_chain_for_region(
                     f"(Metal newComputePipelineState size cap)"
                 )
                 break
+            # BACKWARD-only Metal watchdog cap. Backward fused segments larger
+            # than ``backward_max_segment_nodes`` ops run as ONE grid_chunks
+            # command buffer whose GPU time exceeds the macOS GPU watchdog
+            # (kIOGPUCommandBufferCallbackErrorTimeout) -- the 3-op
+            # ..._chain_7_10 mega-kernel was the FIRST backward segment to die.
+            # Stop extending the backward segment so each backward op lands in
+            # its own command buffer (and the per-row-independent heavy ops can
+            # then be row-windowed below). Forward segments are unaffected.
+            if (
+                execution_phase == DESCRIPTOR_EXECUTION_STAGE_BACKWARD
+                and backward_max_segment_nodes is not None
+                and end - start > backward_max_segment_nodes
+            ):
+                first_failure = (
+                    f"backward direct-buffer segment reached "
+                    f"backward_max_segment_nodes={backward_max_segment_nodes} "
+                    f"(Metal GPU watchdog command-buffer cap)"
+                )
+                break
             target = selector.select(candidate_region)
             if target is None:
                 first_failure = (
@@ -14372,12 +14481,30 @@ def plan_path_c_direct_fusion_chain_for_region(
             # 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.
+            #
+            # ROW-WINDOWING for the per-row-INDEPENDENT heavy backward ops
+            # (sparse_mla_fp8_apply_bwd / attention_qkv_projection_bwd; see
+            # _ROW_CHUNKED_INDEPENDENT_BACKWARD_OPS). At full scale each of these
+            # times out the macOS GPU watchdog even ISOLATED in its own 1-op
+            # grid_chunks command buffer (~10-12s). They have NO carried
+            # reverse-time state (each output row depends only on its own input
+            # row + shared weights), so launcher_chunks splits them into K command
+            # buffers over the independent ROW axis with zero cross-launch state
+            # (_row_phased_launcher_carry_buffers_for_nodes adds carry buffers ONLY
+            # for mamba3/m2rnn, never for these). Weight grads still accumulate
+            # correctly via the path_c_first_row_launch one-time owner-grad zero.
+            # Both recurrent and independent watchdog-heavy ops therefore take the
+            # SAME launcher_chunks treatment; the carry-buffer plumbing distinguishes
+            # them downstream. (Metal-only watchdog split; CUDA keeps grid_chunks.)
             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)
-                    in _TIME_CHUNKED_RECURRENT_BACKWARD_OPS
+                    in (
+                        _TIME_CHUNKED_RECURRENT_BACKWARD_OPS
+                        | _ROW_CHUNKED_INDEPENDENT_BACKWARD_OPS
+                    )
                     for node in candidate_region.nodes
                 )
             ):
@@ -14480,6 +14607,9 @@ def plan_path_c_direct_fusion_chains_for_model(
     forward_max_segment_nodes: int | None | _ResolveFromTarget = (
         _RESOLVE_FORWARD_CAP_FROM_TARGET
     ),
+    backward_max_segment_nodes: int | None | _ResolveFromTarget = (
+        _RESOLVE_BACKWARD_CAP_FROM_TARGET
+    ),
     registry: PathCFusionScheduleRegistry | None = None,
     sequence_length: int | None = None,
 ) -> tuple[PathCFusionScheduleChainPlan, ...]:
@@ -14490,9 +14620,9 @@ def plan_path_c_direct_fusion_chains_for_model(
     model's brick graph first, then split each discovered region into generic
     direct-buffer segments.  It does not consult named acceptance fixtures.
 
-    ``forward_max_segment_nodes`` is forwarded to
-    :func:`plan_path_c_direct_fusion_chain_for_region` (Metal-only forward
-    fusion cap; see that function's docstring).
+    ``forward_max_segment_nodes`` / ``backward_max_segment_nodes`` are forwarded
+    to :func:`plan_path_c_direct_fusion_chain_for_region` (Metal-only fusion caps
+    for the forward / backward phases; see that function's docstring).
     """
 
     regions = build_path_c_model_regions_from_model(
@@ -14509,6 +14639,7 @@ def plan_path_c_direct_fusion_chains_for_model(
             max_kernel_buffers=max_kernel_buffers,
             max_segment_nodes=max_segment_nodes,
             forward_max_segment_nodes=forward_max_segment_nodes,
+            backward_max_segment_nodes=backward_max_segment_nodes,
             registry=registry,
         )
         for region in regions
 
@@ -0,0 +1,127 @@
+#!/usr/bin/env python3
+"""FAST determinism check for row-windowed (launcher_chunks) independent backward.
+
+Runs the FIRST N row-window launches of a per-row-INDEPENDENT backward op TWICE
+on the same random inputs (re-initialising owner grads each time), and asserts
+every DIRECT-WRITE grad buffer is bit-exact across the two runs. The only
+non-deterministic buffer is rope_inv_freq_grad (reduced via relaxed T.atomic_add,
+verified in the generated source) -- its run-to-run float-reassociation variance
+is inherent to the op in ANY dispatch mode and is reported (relative) but not
+required to be bit-exact. This proves the row-windowing + path_c_first_row_launch
+one-time owner-grad zero is deterministic on the non-atomic grads.
+"""
+from __future__ import annotations
+
+import argparse
+import dataclasses
+import sys
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+
+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 S  # noqa: E402
+
+PATH_C_SCALAR_KERNEL_PARAM_DEFAULTS = m.PATH_C_SCALAR_KERNEL_PARAM_DEFAULTS
+path_c_kernel_buffer_order = m.path_c_kernel_buffer_order
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--op", default="attention_qkv_projection_bwd")
+    parser.add_argument("--max-launches", type=int, default=24)
+    args = parser.parse_args()
+
+    profile, route_symbols, regions = m._local_gb10_path_c_model_regions()
+    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)
+    seg = next(s for s in chain.segments if [n.op_name for n in s.region.nodes] == [args.op])
+    seg0 = dataclasses.replace(seg, index=0)
+    chain1 = dataclasses.replace(chain, segments=(seg0,))
+    arts = m.compile_path_c_direct_fusion_chain_artifacts(chain1)
+    art = arts[0]
+    prim_func = seg0.schedule_target.schedule_template(seg0.region)
+
+    specs = m._path_c_direct_chain_required_logical_buffer_specs(chain1)
+    rng = np.random.default_rng(0)
+    init = {}
+    for nm, sp in specs.items():
+        shape = tuple(int(d) for d in sp["shape"]); dt = str(sp["dtype"])
+        if dt in ("float16", "bfloat16", "float32"):
+            arr = rng.standard_normal(shape) * 0.1
+            init[nm] = arr.astype("float32") if dt == "float32" else arr.astype("float16")
+        else:
+            init[nm] = np.zeros(shape, dtype=dt)
+
+    kbuf_shapes = m._path_c_kernel_buffer_shapes(prim_func)
+    korder = path_c_kernel_buffer_order(prim_func)
+    gate = str(getattr(prim_func, "_cppmega_path_c_backward_gate_param", "path_c_run_backward") or "path_c_run_backward")
+    cparam = str(getattr(prim_func, "_cppmega_path_c_row_chunk_index_param", "") or "")
+    sparam = str(getattr(prim_func, "_cppmega_path_c_row_subchunk_index_param", "") or "")
+    launches = m._path_c_segment_time_chunk_launches(prim_func)
+    n = min(len(launches), args.max_launches)
+
+    def run() -> dict[str, np.ndarray]:
+        bufs = {nm: mx.array(init[nm]) for nm in specs}
+        mx.eval(*bufs.values())
+        kargs: list[Any] = []; spos: dict[str, int] = {}
+        for nm in korder:
+            if nm in bufs:
+                kargs.append(m._path_c_exact_kernel_buffer(bufs[nm], kbuf_shapes.get(nm), mx_module=mx))
+            elif nm == gate:
+                kargs.append(1)
+            elif nm in PATH_C_SCALAR_KERNEL_PARAM_DEFAULTS:
+                spos[str(nm)] = len(kargs); kargs.append(PATH_C_SCALAR_KERNEL_PARAM_DEFAULTS[nm])
+            else:
+                raise ValueError(f"unbound {nm}")
+        arrays = tuple(kargs); barrays = tuple(a for a in arrays if hasattr(a, "shape"))
+        mx.eval(*barrays)
+        cp, sp = spos[cparam], spos[sparam]
+        for ci, si in launches[:n]:
+            la = list(arrays); la[cp] = int(ci); la[sp] = int(si)
+            art(*la); mx.eval(*barrays); mx.synchronize()
+        return {nm: np.asarray(bufs[nm].astype(mx.float32)) for nm in bufs}
+
+    # Buffers written via relaxed T.atomic_add are non-deterministic in float
+    # summation order in ANY dispatch mode (dispatch-independent), so they are
+    # reported (relative) but not required bit-exact. Every OTHER grad is a
+    # direct indexed write and MUST be bit-exact across the two row-windowed runs.
+    import re
+    atomic_targets = {
+        mt.group(1)
+        for line in prim_func._cppmega_path_c_generated_source.split("\n")
+        for mt in [re.search(r"T\.atomic_add\(([A-Za-z0-9_]+)", line)]
+        if mt
+    }
+
+    a = run(); b = run()
+    det_max = 0.0; det_worst = None; atomic_max_abs = 0.0; changed = 0
+    for nm in a:
+        if not np.allclose(a[nm], 0):
+            changed += 1
+        if nm in atomic_targets:
+            atomic_max_abs = max(atomic_max_abs, float(np.max(np.abs(a[nm] - b[nm]))) if a[nm].size else 0.0)
+            continue
+        d = float(np.max(np.abs(a[nm] - b[nm]))) if a[nm].size else 0.0
+        if d > det_max:
+            det_max = d; det_worst = nm
+    print(f"[det] op={args.op} launches={n}/{len(launches)} changed_out={changed} "
+          f"atomic_targets={sorted(t.split('_')[-2]+'_'+t.split('_')[-1] for t in atomic_targets)} "
+          f"direct_write_grads_max_abs_diff={det_max:.3e} (worst={det_worst}) "
+          f"atomic_grads_max_abs_diff={atomic_max_abs:.3e}", flush=True)
+    ok = det_max == 0.0 and changed > 0
+    print("[det] PASS (direct-write grads bit-exact; atomic grads vary by float noise)"
+          if ok else "[det] FAIL", flush=True)
+    return 0 if ok else 1
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())