perf(mamba3-bwd): replace lane-disjoint global atomic_add with non-atomic RMW

PROFILE (M4 Max, full-scale local_gb10_quarter, hidden=3584 S=4096):
- GPU-vs-CPU split (Metal System Trace, metal-gpu-intervals): the 0.58-0.64s
  per reverse time-step is ~100% PURE GPU compute. Each launch = ONE Compute
  Command; CPU->GPU latency ~0.2-1.0ms (<0.2%). Not a CPU/eval/sync problem.
- Dispatch: T.Kernel(1, threads=1024) -> a SINGLE threadgroup => ~1/40 of the
  40 GPU cores. Confirmed from the generated PrimFunc source.
- Replay fraction: checkpoint-interval sweep {8,4,2,1} -> steady 0.64/0.54/
  0.52/0.510s. Even with ZERO replay (interval=1) it is 0.510s/step, so the
  per-step backward math (not the replay) dominates the floor; each extra
  replay row adds only ~0.039s.
- Limiter: NOT ALU. Diagnostic (atomic->scratch) showed the in_proj_weight_grad
  scatter loop's entire 0.198s/step cost IS the two relaxed global atomic_add
  ops -- removing the atomics but keeping the 18784xH loop+MAC was free. Metal
  serializes relaxed global atomics; the kernel is atomic/memory-bound.

FIX: every grad-scatter whose global address is uniquely owned by its thread
lane within a launch (in/out-proj weight, conv weight/bias, dt_bias, B/C norm
& bias) is converted from T.atomic_add to a plain read-modify-write. The time
loop is serial within the single threadgroup, so cross-step accumulation stays
correct; no two lanes ever touch the same address in one step. hidden_grad and
the shared-scratch reductions (D, a/dt_grad, b/c_group_grad, project_grad angle)
genuinely collide across lanes and KEEP their atomics.

RESULT (full scale):
- no-replay floor: 0.510 -> 0.348s/step (1.47x, -32%)
- production interval=8 block average: ~0.595 -> ~0.466s/step (~1.28x, -22%)
- window=2 steady: ~1.19s -> ~0.94s (~0.47s/step)

GRADIENT PARITY: bitwise-identical. tiny-smoke fwd+bwd, 15 non-zero mamba3 grad
buffers, worst max|baseline-fixed| = 0.000e+00.

Adds isolation probes: _bwd_shape_introspect (analytic per-step cost),
_bwd_dump_source (kernel launch + loop structure), _bwd_ckpt_sweep (replay
fraction). The dominant remaining inefficiency is the single-threadgroup
dispatch (~1/40 GPU); see report for the grid-parallel feasibility plan.

Author: apstenku123

cppmega_mlx/runtime/path_c_fusion_schedules.py

@@ -11581,10 +11581,12 @@ def _mamba3_emit_recompute_row(
             access_by_buffer,
             f"{out_dim} * {inner_dim} + {feature}",
         )
+        # Lane-disjoint: lane==feature, address==out_dim*inner_dim+feature, so
+        # each lane owns a distinct out_proj_weight_grad column. Serial time loop
+        # => non-atomic RMW is byte-identical and avoids the global-atomic stall.
         body.append(
-            f"{indent * 8}T.atomic_add({out_grad_ref}, "
-            f"{stage_grad}[0] * {out_inner}[{feature}], "
-            "memory_order=\"relaxed\")"
+            f"{indent * 8}{out_grad_ref} = {out_grad_ref} + "
+            f"({stage_grad}[0] * {out_inner}[{feature}])"
         )
     body.append(f"{indent * 6}T.sync_threads()")
     body.append(
@@ -11845,17 +11847,17 @@ def _mamba3_emit_recompute_row(
             access_by_buffer,
             grad_flat,
         )
+        # Lane-disjoint: lane==rank_loop, grad_flat=(rank_loop*groups+.)*sd+.
+        # so each lane owns distinct addresses; serial time loop => RMW.
         body.append(
-            f"{indent * 9}T.atomic_add({b_norm_grad_ref}, "
-            f"{scalar0}[0] * {conv}[{conv_b_offset} + {grad_flat}] * "
-            f"{b_inv}[{rank_loop} * {groups} + {group_loop}], "
-            "memory_order=\"relaxed\")"
+            f"{indent * 9}{b_norm_grad_ref} = {b_norm_grad_ref} + "
+            f"({scalar0}[0] * {conv}[{conv_b_offset} + {grad_flat}] * "
+            f"{b_inv}[{rank_loop} * {groups} + {group_loop}])"
         )
     if b_bias_grad is not None:
         b_bias_grad_ref = _indexed_buffer_ref(b_bias_grad, access_by_buffer, grad_flat)
         body.append(
-            f"{indent * 9}T.atomic_add({b_bias_grad_ref}, {scalar0}[0], "
-            "memory_order=\"relaxed\")"
+            f"{indent * 9}{b_bias_grad_ref} = {b_bias_grad_ref} + {scalar0}[0]"
         )
     body.append(
         f"{indent * 9}{conv_grad}[{conv_b_offset} + {grad_flat}] = "
@@ -11876,17 +11878,16 @@ def _mamba3_emit_recompute_row(
             access_by_buffer,
             grad_flat,
         )
+        # Lane-disjoint (lane==rank_loop); serial time loop => non-atomic RMW.
         body.append(
-            f"{indent * 9}T.atomic_add({c_norm_grad_ref}, "
-            f"{scalar1}[0] * {conv}[{conv_c_offset} + {grad_flat}] * "
-            f"{c_inv}[{rank_loop} * {groups} + {group_loop}], "
-            "memory_order=\"relaxed\")"
+            f"{indent * 9}{c_norm_grad_ref} = {c_norm_grad_ref} + "
+            f"({scalar1}[0] * {conv}[{conv_c_offset} + {grad_flat}] * "
+            f"{c_inv}[{rank_loop} * {groups} + {group_loop}])"
         )
     if c_bias_grad is not None:
         c_bias_grad_ref = _indexed_buffer_ref(c_bias_grad, access_by_buffer, grad_flat)
         body.append(
-            f"{indent * 9}T.atomic_add({c_bias_grad_ref}, {scalar1}[0], "
-            "memory_order=\"relaxed\")"
+            f"{indent * 9}{c_bias_grad_ref} = {c_bias_grad_ref} + {scalar1}[0]"
         )
     body.append(
         f"{indent * 9}{conv_grad}[{conv_c_offset} + {grad_flat}] = "
@@ -11935,9 +11936,9 @@ def _mamba3_emit_recompute_row(
     )
     if dt_bias_grad is not None:
         dt_bias_grad_ref = _indexed_buffer_ref(dt_bias_grad, access_by_buffer, hidden_dim)
+        # Lane-disjoint (lane==hidden_dim); serial time loop => non-atomic RMW.
         body.append(
-            f"{indent * 9}T.atomic_add({dt_bias_grad_ref}, {scalar3}[0], "
-            "memory_order=\"relaxed\")"
+            f"{indent * 9}{dt_bias_grad_ref} = {dt_bias_grad_ref} + {scalar3}[0]"
         )
     body.append(f"{indent * 9}for {hidden_loop} in T.serial(0, {hidden_size}):")
     if in_proj_weight_grad is not None:
@@ -11951,15 +11952,16 @@ def _mamba3_emit_recompute_row(
             access_by_buffer,
             f"({trap_offset} + {hidden_dim}) * {hidden_size} + {hidden_loop}",
         )
+        # Lane-disjoint (lane==hidden_dim => distinct rows dt_offset+hidden_dim /
+        # trap_offset+hidden_dim). The later main scatter touches the same rows
+        # only across a sync_threads barrier (sequential) => non-atomic RMW.
         body.append(
-            f"{indent * 10}T.atomic_add({next_dt_grad_ref}, "
-            f"{_mamba3_hidden_expr(f'({time_idx} + 1)', hidden_loop)} * {scalar3}[0], "
-            "memory_order=\"relaxed\")"
+            f"{indent * 10}{next_dt_grad_ref} = {next_dt_grad_ref} + "
+            f"({_mamba3_hidden_expr(f'({time_idx} + 1)', hidden_loop)} * {scalar3}[0])"
         )
         body.append(
-            f"{indent * 10}T.atomic_add({next_trap_grad_ref}, "
-            f"{_mamba3_hidden_expr(f'({time_idx} + 1)', hidden_loop)} * {scalar4}[0], "
-            "memory_order=\"relaxed\")"
+            f"{indent * 10}{next_trap_grad_ref} = {next_trap_grad_ref} + "
+            f"({_mamba3_hidden_expr(f'({time_idx} + 1)', hidden_loop)} * {scalar4}[0])"
         )
     if hidden_grad is not None:
         hidden_grad_ref = _indexed_buffer_ref(
@@ -12012,9 +12014,11 @@ def _mamba3_emit_recompute_row(
     )
     if dt_bias_grad is not None:
         dt_bias_grad_ref = _indexed_buffer_ref(dt_bias_grad, access_by_buffer, head)
+        # Lane-disjoint (lane==head); separated from the earlier dt_bias scatter
+        # by a sync_threads => non-atomic RMW accumulates correctly.
         body.append(
-            f"{indent * 7}T.atomic_add({dt_bias_grad_ref}, "
-            f"{dt_grad}[{head}] * {scalar0}[0], memory_order=\"relaxed\")"
+            f"{indent * 7}{dt_bias_grad_ref} = {dt_bias_grad_ref} + "
+            f"({dt_grad}[{head}] * {scalar0}[0])"
         )
     body.append(
         f"{indent * 7}if -T.log(1.0 + T.exp({projected}[{a_offset} + {head}])) < -0.01:"
@@ -12038,9 +12042,9 @@ def _mamba3_emit_recompute_row(
     )
     if conv_bias_grad is not None:
         conv_bias_ref = _indexed_buffer_ref(conv_bias_grad, access_by_buffer, conv_ch)
+        # Lane-disjoint (lane==conv_ch); serial time loop => non-atomic RMW.
         body.append(
-            f"{indent * 7}T.atomic_add({conv_bias_ref}, {scalar1}[0], "
-            "memory_order=\"relaxed\")"
+            f"{indent * 7}{conv_bias_ref} = {conv_bias_ref} + {scalar1}[0]"
         )
     if history_len > 0:
         body.append(f"{indent * 7}for {kernel_pos} in T.serial(0, {history_len}):")
@@ -12072,9 +12076,11 @@ def _mamba3_emit_recompute_row(
                 access_by_buffer,
                 f"{conv_ch} * {kernel} + {kernel_pos}",
             )
+            # Lane-disjoint: lane==conv_ch owns row conv_ch*kernel+. ; serial time
+            # loop => non-atomic RMW byte-identical to the atomic.
             body.append(
-                f"{indent * 8}T.atomic_add({conv_weight_grad_ref}, "
-                f"{scalar1}[0] * {scalar2}[0], memory_order=\"relaxed\")"
+                f"{indent * 8}{conv_weight_grad_ref} = {conv_weight_grad_ref} + "
+                f"({scalar1}[0] * {scalar2}[0])"
             )
         body.append(f"{indent * 8}if {src_row} >= 0:")
         body.append(f"{indent * 9}for {hidden_loop} in T.serial(0, {hidden_size}):")
@@ -12084,11 +12090,13 @@ def _mamba3_emit_recompute_row(
                 access_by_buffer,
                 f"({x_offset} + {conv_ch}) * {hidden_size} + {hidden_loop}",
             )
+            # Lane-disjoint: lane==conv_ch => row (x_offset+conv_ch)*H+. ; the
+            # main in-proj scatter (a sync_threads later) hits the same rows only
+            # SEQUENTIALLY, so non-atomic RMW accumulates correctly.
             body.append(
-                f"{indent * 10}T.atomic_add({in_proj_grad_ref}, "
-                f"{_mamba3_hidden_expr(src_row, hidden_loop)} * {scalar1}[0] * "
-                f"{_mamba3_conv_weight_expr(conv_ch, kernel_pos)}, "
-                "memory_order=\"relaxed\")"
+                f"{indent * 10}{in_proj_grad_ref} = {in_proj_grad_ref} + "
+                f"({_mamba3_hidden_expr(src_row, hidden_loop)} * {scalar1}[0] * "
+                f"{_mamba3_conv_weight_expr(conv_ch, kernel_pos)})"
             )
         if hidden_grad is not None:
             hidden_grad_ref = _indexed_buffer_ref(
@@ -12108,10 +12116,11 @@ def _mamba3_emit_recompute_row(
             access_by_buffer,
             f"{conv_ch} * {kernel} + {history_len}",
         )
+        # Lane-disjoint (lane==conv_ch); serial time loop => non-atomic RMW.
         body.append(
-            f"{indent * 7}T.atomic_add({current_conv_weight_grad_ref}, "
-            f"{scalar1}[0] * {projected}[{x_offset} + {conv_ch}], "
-            "memory_order=\"relaxed\")"
+            f"{indent * 7}{current_conv_weight_grad_ref} = "
+            f"{current_conv_weight_grad_ref} + "
+            f"({scalar1}[0] * {projected}[{x_offset} + {conv_ch}])"
         )
     body.append(
         f"{indent * 7}{project_grad}[{x_offset} + {conv_ch}] = "
@@ -12130,12 +12139,18 @@ def _mamba3_emit_recompute_row(
             access_by_buffer,
             f"{proj_dim} * {hidden_size} + {hidden_loop}",
         )
+        # Lane-disjoint scatter: lane==proj_dim, so each lane owns a distinct
+        # in_proj_weight_grad ROW (proj_dim*H + .) and the serial time loop runs
+        # one step at a time. No two threads ever target the same address within
+        # a launch -> a non-atomic read-modify-write is byte-identical to the
+        # atomic and ~10x cheaper (relaxed global atomics serialize on Metal).
         body.append(
-            f"{indent * 8}T.atomic_add({in_proj_grad_ref}, "
-            f"{_mamba3_hidden_expr(time_idx, hidden_loop)} * {project_grad}[{proj_dim}], "
-            "memory_order=\"relaxed\")"
+            f"{indent * 8}{in_proj_grad_ref} = {in_proj_grad_ref} + "
+            f"({_mamba3_hidden_expr(time_idx, hidden_loop)} * {project_grad}[{proj_dim}])"
         )
     if hidden_grad is not None:
+        # hidden_grad[time_idx*H + hidden_loop]: ALL proj_dim lanes accumulate
+        # into the same hidden_loop column -> lanes COLLIDE; keep atomic.
         hidden_grad_ref = _indexed_buffer_ref(
             hidden_grad,
             access_by_buffer,

scripts/_bwd_ckpt_sweep.py

@@ -0,0 +1,125 @@
+#!/usr/bin/env python3
+"""Replay-fraction probe: steady per-launch time vs checkpoint interval.
+
+Drives the SAME mamba3 bwd window=1 launches as the window sweep, but for a
+given ``--checkpoint-interval`` (monkeypatched). Because window=1 fixes one
+reverse time-step per launch and the replay loop re-runs the forward recompute
+for ``[checkpoint_start, time_idx]`` (up to checkpoint_interval rows), the steady
+per-launch time isolates the REPLAY cost. Smaller interval => fewer replays =>
+less recompute => faster, IF replay dominates.
+
+We time the FIRST launch (time_idx=4095) where the replay length == the worst
+case for that interval, plus the steady launches. Reports per-launch time so we
+can attribute the fraction that is checkpoint-replay vs the fixed backward math.
+"""
+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):
+    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")
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--checkpoint-interval", type=int, required=True)
+    parser.add_argument("--max-launches", type=int, default=5)
+    args = parser.parse_args()
+
+    sched.MAMBA3_BWD_ROWS_PER_KERNEL_LAUNCH = 1
+    sched.MAMBA3_BWD_REPLAY_CHECKPOINT_INTERVAL = int(args.checkpoint_interval)
+    print(f"[ckpt] checkpoint_interval={args.checkpoint_interval} window=1", flush=True)
+
+    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 = _find_mamba3_segment(chain)
+    artifacts = m.compile_path_c_direct_fusion_chain_artifacts(chain)
+    sub_chain = dataclasses.replace(chain, segments=(mamba_seg,))
+
+    launch_times: list[float] = []
+    state: dict[str, Any] = {"last": None}
+    real_sync = mx.synchronize
+
+    def timed_sync(*a, **k):
+        r = real_sync(*a, **k)
+        now = time.perf_counter()
+        if state["last"] is not None:
+            launch_times.append(now - state["last"])
+        state["last"] = now
+        return r
+
+    mx.synchronize = timed_sync
+
+    real_launches_fn = m._path_c_segment_time_chunk_launches
+    limit = int(args.max_launches)
+
+    def limited_launches(pf):
+        full = real_launches_fn(pf)
+        return full[:limit] if full else full
+
+    m._path_c_segment_time_chunk_launches = limited_launches
+
+    specs = m._path_c_direct_chain_required_logical_buffer_specs(sub_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())
+
+    status = "ok"
+    err = None
+    t0 = time.perf_counter()
+    state["last"] = t0
+    try:
+        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()
+    finally:
+        mx.synchronize = real_sync
+        m._path_c_segment_time_chunk_launches = real_launches_fn
+
+    for i, dt in enumerate(launch_times):
+        print(f"[ckpt]   launch[{i}] = {dt:.3f}s", flush=True)
+    steady = (sum(launch_times[1:]) / max(1, len(launch_times) - 1)
+              if len(launch_times) > 1 else None)
+    out = {
+        "checkpoint_interval": int(args.checkpoint_interval),
+        "status": status,
+        "error": err,
+        "first_launch_s": round(launch_times[0], 4) if launch_times else None,
+        "steady_per_launch_s": round(steady, 4) if steady is not None else None,
+    }
+    print("\n[ckpt] RESULT:\n" + json.dumps(out, indent=2), flush=True)
+    return 0 if status == "ok" else 7
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())