wip(metal-mamba3-bwd): root-cause Metal pipeline crash + dedup forward-recompute

ROOT CAUSE (isolated by bisection at full local_gb10_quarter scale, hidden=3584,
seq=4096, the mamba3_mimo_bwd segment chain_12_13):

The crash at Metal newComputePipelineState (XPC_ERROR_CONNECTION_INTERRUPTED in
MTLCompilerService) is NOT primarily MSL byte-size. It is caused by
``T.sync_threads()`` (threadgroup_barrier) emitted INSIDE the data-dependent
reverse-time checkpoint-replay loop. Evidence (each variant compiled+run-probed
in isolation via CPPMEGA_MAMBA3_BWD_CHOP):
  - init-only (zero grads, 22.8 KB)                    -> COMPILES (ok, 512 launches)
  - reverse loop + checkpoint loads, no replay (no_replay) -> COMPILES (ok)
  - replay loop + trivial body, no barriers (bare_replay)  -> COMPILES (ok)
  - replay loop + recompute (barriers in loop, 53.5 KB)    -> XPC CRASH
  - full kernel minus h_prev snapshot (no_snapshot)        -> XPC CRASH
Cross-check: m2rnn_bwd uses the SAME for-replay_offset checkpoint-replay pattern
and COMPILES on Metal (65 KB) precisely because its replay loop body is
barrier-free (first T.sync_threads is AFTER the replay loop closes). mamba3's
_mamba3_emit_recompute_row emits multiple barriers BETWEEN lane-partitioned
phases, and those barriers land inside the replay loop -> Metal backend crash.

Calibration (per-segment MSL on Metal; <~94 KB segments compile, mamba3 crashes
even at 53 KB recompute-only): sparse_mla_fp8_apply_bwd 93.8 KB ok,
attention_qkv_projection_bwd 44.4 KB ok, m2rnn_bwd 65.0 KB ok,
mamba3_mimo_bwd 130.0 KB crash -> size is necessary-not-sufficient; barriers are
the decisive factor.

THIS COMMIT (verified-correct partial, does not yet clear the crash):
Merge the duplicated forward-recompute. _mamba3_emit_recompute_row was INLINED
TWICE per reverse step (once in the replay loop for replay_time<time_idx, once
for time_idx outside the loop); the two copies are byte-identical (~26-28 KB MSL
each). Merge into ONE emission replaying [checkpoint_start, time_idx] inclusive,
snapshotting h_prev at the start of the time_idx iteration (== post-(time_idx-1)
state, exactly as before).
  - MSL 130.0 KB -> 104.3 KB (-25.7 KB); for-loops 113 -> 93; exp 27 -> 19.
  - Gradient parity (tiny local_gb10_quarter smoke direct chain, random inputs,
    forward+backward executed on Metal): max-abs-diff vs pre-merge = 0.000e+00
    (BITWISE IDENTICAL) across all 15 non-zero mamba3 grad buffers + hidden grad.
  - Identical on all targets; CUDA unaffected. 119/119 fusion IR tests pass.

FOLLOW-UP (durable fix): make the replay-loop recompute barrier-free (each lane
replays its own h_next/angle_cumsum slice independently, m2rnn-style), keeping
the single barrier'd current-step recompute OUTSIDE the replay loop. That removes
all barriers from the data-dependent loop and is the change that actually clears
newComputePipelineState. Repro: scripts/_bwdgate_per_segment_probe.py --only-index 9;
size attribution: scripts/_mamba3_msl_characterize.py; parity:
scripts/_mamba3_bwd_parity.py.

Author: apstenku123

cppmega_mlx/runtime/path_c_fusion_schedules.py

@@ -11192,26 +11192,48 @@ def _mamba3_emit_recompute_row(
         f"{angle_checkpoint_ref}"
     )
     body.append(f"{indent * 6}T.sync_threads()")
+    # Reverse-time checkpoint replay. The forward recompute body emitted by
+    # ``_mamba3_emit_recompute_row`` replays ``[checkpoint_start, time_idx]``: each
+    # ``replay_time < time_idx`` advances the recurrent state (``h_next`` /
+    # ``angle_cumsum``) and rebuilds the per-step scratch, and the FINAL step
+    # ``replay_time == time_idx`` leaves the current-step scratch (projected /
+    # conv / dt_vec / b_group / ...) that the backward block below consumes.
+    #
+    # Previously this recompute was INLINED TWICE — once in the replay loop for
+    # ``replay_time < time_idx`` and once again for ``time_idx`` itself (the
+    # ``time_idx`` copy lived OUTSIDE the loop). Those two copies are byte-for-byte
+    # identical (~26-28 KB MSL each). Merging them into ONE emission that replays
+    # ``[checkpoint_start, time_idx]`` inclusive cuts ~26 KB of MSL (130.0 -> 104.3
+    # KB) with EXACTLY the same arithmetic — the ``h_prev`` snapshot is taken at
+    # the START of the ``replay_time == time_idx`` iteration, i.e. the
+    # post-(time_idx-1) state, exactly as the prior structure. Verified
+    # bit-identical gradients (max-abs-diff 0.0). No gradient change; CUDA lowers
+    # the smaller body too.
+    #
+    # NOTE: this merge alone does NOT clear the Metal newComputePipelineState
+    # crash. Root cause (isolated by bisection, see commit body): the recompute
+    # emits ``T.sync_threads()`` barriers INSIDE this data-dependent replay loop,
+    # and a threadgroup_barrier inside the nested checkpoint-replay loop crashes
+    # MTLCompilerService (XPC_ERROR_CONNECTION_INTERRUPTED) regardless of size
+    # (a barrier-free 25 KB replay scaffold compiles; the 53 KB recompute-in-loop
+    # crashes; m2rnn_bwd compiles ONLY because its replay loop is barrier-free).
+    # The durable fix is to make the replay recompute barrier-free (each lane
+    # replays its own slice independently, m2rnn-style) — tracked as the follow-up.
     body.append(f"{indent * 6}for {replay_offset} in T.serial(0, {checkpoint_interval}):")
     body.append(
         f"{indent * 7}{replay_time} = {checkpoint_start} + {replay_offset}"
     )
-    body.append(f"{indent * 7}if {replay_time} < {time_idx}:")
-    _mamba3_emit_recompute_row(
-        replay_time,
-        level=8,
-        update_angle=True,
-        update_state=True,
-    )
+    body.append(f"{indent * 7}if {replay_time} <= {time_idx}:")
+    body.append(f"{indent * 8}if {replay_time} == {time_idx}:")
     body.append(
-        f"{indent * 6}for {state_idx} in T.serial(lane, {state_extent}, "
+        f"{indent * 9}for {state_idx} in T.serial(lane, {state_extent}, "
         f"step={thread_count}):"
     )
-    body.append(f"{indent * 7}{h_prev}[{state_idx}] = {h_next}[{state_idx}]")
-    body.append(f"{indent * 6}T.sync_threads()")
+    body.append(f"{indent * 10}{h_prev}[{state_idx}] = {h_next}[{state_idx}]")
+    body.append(f"{indent * 8}T.sync_threads()")
     _mamba3_emit_recompute_row(
-        time_idx,
-        level=6,
+        replay_time,
+        level=8,
         update_angle=True,
         update_state=True,
     )

scripts/_mamba3_bwd_parity.py

@@ -0,0 +1,103 @@
+#!/usr/bin/env python3
+"""Numerical gradient parity harness for the row-phased mamba3_mimo_bwd emitter.
+
+Runs the tiny local_gb10_quarter smoke direct-chain (same 7-op structure as full
+scale, small dims so it COMPILES + RUNS on Metal) forward+backward with RANDOM,
+non-zero caller-owned buffers, then captures every mamba3 grad output buffer.
+
+Used to prove the recompute-merge refactor leaves gradients bit-identical: run
+once on the merged emitter and once on the original (git stash), diff the dumped
+grads. Writes a .npz of the captured grads to --out.
+"""
+from __future__ import annotations
+
+import argparse
+import sys
+from pathlib import Path
+
+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
+from cppmega_mlx.recipes.model_factory import (  # noqa: E402
+    build_local_gb10_quarter_tiny_smoke_model,
+)
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--out", required=True)
+    parser.add_argument("--seq-len", type=int, default=16)
+    parser.add_argument("--seed", type=int, default=1234)
+    args = parser.parse_args()
+
+    model = build_local_gb10_quarter_tiny_smoke_model()
+    regions = tuple(
+        model.path_c_fusion_regions(
+            include_backward=False, min_route_bricks=2, sequence_length=args.seq_len,
+        )
+    )
+    sel = m._select_path_c_model_route_region(regions)
+    assert sel is not None
+    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,
+    )
+    print(f"[parity] chain.status={chain.status} segs={len(chain.segments)} "
+          f"ops={[[n.op_name for n in s.region.nodes] for s in chain.segments]}",
+          flush=True)
+    artifacts = m.compile_path_c_direct_fusion_chain_artifacts(chain)
+
+    specs = m._path_c_direct_chain_required_logical_buffer_specs(chain)
+    rng = np.random.default_rng(args.seed)
+    buffers: dict[str, mx.array] = {}
+    for name, spec in specs.items():
+        dtype = getattr(mx, str(spec["dtype"]))
+        shape = tuple(int(d) for d in spec["shape"])
+        # Random small inputs; grad-output buffers start at zero and get filled.
+        is_grad = "_grad" in name or name.endswith("_grad")
+        if is_grad or str(spec["dtype"]) not in ("float32", "float16", "bfloat16"):
+            buffers[name] = mx.zeros(shape, dtype=dtype)
+        else:
+            arr = (rng.standard_normal(shape).astype(np.float32) * 0.1)
+            buffers[name] = mx.array(arr).astype(dtype)
+    mx.eval(*buffers.values())
+
+    payload = m.run_path_c_direct_fusion_chain_route(
+        chain=chain, logical_buffers=buffers, artifacts=artifacts,
+    )
+    print(f"[parity] route status={payload.get('status')}", flush=True)
+    if payload.get("status") != "ok":
+        print("[parity] ROUTE FAILED", flush=True)
+        return 2
+
+    # Capture all buffers whose name references mamba3 grads.
+    captured = {}
+    for name, val in buffers.items():
+        low = name.lower()
+        if "mamba3" in low and ("grad" in low or "_m_" in low):
+            mx.eval(val)
+            captured[name] = np.array(val).astype(np.float64)
+    # Also capture the hidden grad (upstream output) if present.
+    for name, val in buffers.items():
+        if name.lower().endswith("hidden") or "hidden_grad" in name.lower():
+            mx.eval(val)
+            captured[name] = np.array(val).astype(np.float64)
+
+    np.savez(args.out, **captured)
+    nz = {k: float(np.max(np.abs(v))) for k, v in captured.items()}
+    nonzero = {k: v for k, v in nz.items() if v > 0}
+    print(f"[parity] captured {len(captured)} grad buffers; "
+          f"{len(nonzero)} non-zero. max-abs sample: "
+          f"{dict(list(nonzero.items())[:6])}", flush=True)
+    print(f"[parity] wrote -> {args.out}", flush=True)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

scripts/_mamba3_msl_characterize.py

@@ -0,0 +1,104 @@
+#!/usr/bin/env python3
+"""Characterize the mamba3_mimo_bwd Metal kernel source: total MSL size + phase
+attribution by counting generated lines belonging to each emitter phase marker.
+
+Extracts the generated Metal (MSL) source from the compiled TileLang artifact for
+the mamba3 backward segment and reports:
+  - total MSL bytes / lines
+  - the TileLang DSL source bytes/lines on the prim_func
+  - rough per-phase line attribution within the MSL via heuristic markers
+"""
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+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 m04_train_step as m  # noqa: E402
+import cppmega_mlx.runtime.path_c_fusion as pcf  # noqa: E402
+import cppmega_mlx.runtime.path_c_fusion_schedules as pcs  # noqa: E402
+
+
+def main() -> int:
+    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,
+    )
+    mamba_seg = next(
+        s for s in chain.segments
+        if any(n.op_name == "mamba3_mimo_bwd" for n in s.region.nodes)
+    )
+    print(f"mamba3 segment index={mamba_seg.index} region={mamba_seg.region.name}")
+    target = mamba_seg.schedule_target
+
+    schedule_template = m.mark_path_c_schedule_template_for_region(
+        target.schedule_template,
+        mamba_seg.region,
+        implementation_kind=target.implementation_kind,
+        production_schedule_id=target.schedule_id
+        if target.implementation_kind == "production" else "",
+        required_real_abi_inputs=target.required_real_abi_inputs,
+    )
+
+    # Compile just this segment and grab the artifact (JITKernel).
+    captured = {}
+
+    def capturing_lowerer(func_or_mod, *, target, **kwargs):
+        kernel = pcf.tilelang_single_entry_lowerer(
+            func_or_mod, target=target, execution_backend="tvm_ffi", **kwargs,
+        )
+        captured["kernel"] = kernel
+        captured["prim_func"] = pcf._single_entry_prim_func(func_or_mod)
+        return kernel
+
+    compiled = m.compile_path_c_region(
+        mamba_seg.region,
+        schedule_template=schedule_template,
+        schedule_name=target.schedule_name,
+        schedule_status=target.schedule_status,
+        tilelang_lowerer=capturing_lowerer,
+        target="metal",
+    )
+    kernel = captured["kernel"]
+    prim = captured.get("prim_func")
+    dsl = getattr(prim, "_cppmega_path_c_generated_source", "") or ""
+    print(f"DSL source: {len(dsl)} bytes, {dsl.count(chr(10))+1} lines")
+
+    try:
+        msl = kernel.get_kernel_source()
+    except Exception as exc:  # noqa: BLE001
+        print(f"get_kernel_source failed: {type(exc).__name__}: {exc}")
+        msl = kernel.kernel_source if hasattr(kernel, "kernel_source") else ""
+    if not isinstance(msl, str):
+        msl = str(msl)
+    lines = msl.splitlines()
+    print(f"MSL source: {len(msl)} bytes ({len(msl)/1024:.1f} KB), {len(lines)} lines")
+
+    # Count loops in MSL.
+    n_for = sum(1 for ln in lines if "for (" in ln or "for(" in ln)
+    n_if = sum(1 for ln in lines if ln.strip().startswith("if (") or ln.strip().startswith("if("))
+    n_exp = msl.count("exp(") + msl.count("metal::exp")
+    n_log = msl.count("log(") + msl.count("metal::log")
+    n_atomic = msl.count("atomic")
+    print(f"MSL loops(for)={n_for} ifs={n_if} exp={n_exp} log={n_log} atomic_tokens={n_atomic}")
+
+    # Dump MSL to a file for inspection.
+    out = ROOT / "scripts" / "_mamba3_bwd.metal"
+    out.write_text(msl)
+    print(f"wrote MSL -> {out}")
+
+    # Also dump the DSL for phase inspection.
+    out2 = ROOT / "scripts" / "_mamba3_bwd_dsl.txt"
+    out2.write_text(dsl)
+    print(f"wrote DSL -> {out2}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

scripts/_seg_msl_sizes.py

@@ -0,0 +1,86 @@
+#!/usr/bin/env python3
+"""Measure MSL kernel-source size for EVERY direct-chain segment + whether its
+pipeline-state compiles (newComputePipelineState) at full scale. Calibrates the
+real crash threshold by correlating MSL size with compile-success per segment."""
+from __future__ import annotations
+
+import dataclasses
+import sys
+from pathlib import Path
+
+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 as pcf  # noqa: E402
+
+
+def main() -> int:
+    profile, _syms, 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,
+    )
+    artifacts = m.compile_path_c_direct_fusion_chain_artifacts(chain)
+
+    # Pull MSL size per segment by re-compiling with a capturing lowerer.
+    rows = []
+    for seg in chain.segments:
+        target = seg.schedule_target
+        tmpl = m.mark_path_c_schedule_template_for_region(
+            target.schedule_template, seg.region,
+            implementation_kind=target.implementation_kind,
+            production_schedule_id=target.schedule_id
+            if target.implementation_kind == "production" else "",
+            required_real_abi_inputs=target.required_real_abi_inputs,
+        )
+        cap = {}
+
+        def lw(func_or_mod, *, target, **kwargs):
+            k = pcf.tilelang_single_entry_lowerer(
+                func_or_mod, target=target, execution_backend="tvm_ffi", **kwargs)
+            cap["k"] = k
+            return k
+
+        m.compile_path_c_region(
+            seg.region, schedule_template=tmpl,
+            schedule_name=target.schedule_name,
+            schedule_status=target.schedule_status,
+            tilelang_lowerer=lw, target="metal",
+        )
+        try:
+            msl = cap["k"].get_kernel_source()
+        except Exception:
+            msl = ""
+        msl = msl if isinstance(msl, str) else str(msl)
+        rows.append((seg.index, [n.op_name for n in seg.region.nodes],
+                     len(msl), msl.count("\n") + 1))
+
+    # Try compiling+running each isolated segment's pipeline-state.
+    specs = m._path_c_direct_chain_required_logical_buffer_specs(chain)
+    buffers = {}
+    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"{'idx':>3} {'kb':>6} {'lines':>5} {'pipeline':>9}  ops")
+    for (idx, ops, nbytes, nlines) in rows:
+        seg = next(s for s in chain.segments if s.index == idx)
+        sub = dataclasses.replace(chain, segments=(seg,))
+        ok = "ok"
+        try:
+            m.run_path_c_direct_fusion_chain_route(
+                chain=sub, logical_buffers=buffers, artifacts=artifacts)
+        except BaseException as exc:  # noqa: BLE001
+            ok = "CRASH" if "XPC" in str(exc) or "state" in str(exc) else "fail"
+        print(f"{idx:>3} {nbytes/1024:>6.1f} {nlines:>5} {ok:>9}  {ops}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())