feat(mamba3-pathc): Stage 3 — LIVE region-build backward 1->3 surface flip (B2/B1/B0)

Mirrors the forward Stage-2 region flip for the backward (design §2/§3.1/§7):
- _emit_mamba3_chunked_bwd_model_brick_surfaces emits the 3 _bwd surfaces
  (B2->B1->B0, the transpose of F2->F1->F0) when the flag is ON, wired by the
  per-brick grad-handoff buffers (dchunk_states B2->B1, dstates B1->B0,
  dinp_diag/dA_cumsum_y B2->B0, dA_cumsum_tail B1->B0). REUSES the forward-
  materialized cb/dA_cumsum/prev_states (no 8x checkpoint-replay). Invoked from
  the flag-ON _emit_mamba3_chunked_model_brick_surfaces (the forward owner_output
  surfaces carry no AOT autograd, so the backward is emitted explicitly).
- 6 backward grad-handoff buffers registered: _buffer_shape (fp32 region shapes),
  _buffer_dtype (fp32 recurrent-grad precision), and the force-spill set
  DESCRIPTOR_MAMBA3_CHUNKED_BWD_HANDOFF_ABI_BUFFERS.
- FIX: _path_c_schedule_node_execution_phase classified ANY op in the (now-shared)
  delegation set as forward; excluded *_bwd so the 3 backward ops classify as
  BACKWARD phase. The per-node backward stage planner already isolates each into
  its own single-node stage (interpose requires len(nodes)==1).

Verified (tests/test_mamba3_chunked_backward_b0b1b2.py, 11 tests):
- flag OFF: no chunked _bwd surfaces, serial mamba3_mimo unchanged (merge-safe).
- flag ON: 3 _bwd surfaces, all backward-phase, each own backward stage, correct
  B2->B1->B0 grad-handoff + prev_states reuse, fp32 force-spilled handoff buffers.
- region-build flip inspection: 7 nodes (entry_rmsnorm + 3 fwd + 3 bwd), stages
  g0..g3 forward / b0..b2 backward single-node each.

Flag-OFF unchanged: 151 forward + path_c template/IR tests green (1 Stage-2 test
updated to count forward-only chunked nodes). Flag default OFF (merge-safe).

Author: apstenku123

cppmega_mlx/runtime/path_c_fusion.py

@@ -1936,6 +1936,25 @@ def _path_c_model_brick_surface_lowerer_for(
     "prev_states",
     "final_state",
 )
+# Caller-owned BACKWARD grad-handoff buffer suffixes that wire B2 -> B1 -> B0
+# (Stage 3; the analytic transpose of the forward F2 -> F1 -> F0 handoff). They
+# are NEW logical region buffers (registered in _buffer_shape / _buffer_dtype and
+# the schedules force-spill set DESCRIPTOR_MAMBA3_CHUNKED_BWD_HANDOFF_ABI_BUFFERS):
+#   dchunk_states : B2 -> B1   grad of prev_states (per-chunk entry state)
+#   dstates       : B1 -> B0   grad of summary_states
+#   dinp_diag     : B2 -> B0   diag input-outer grad
+#   dA_cumsum_y   : B2 -> B0   Y-path dA_cumsum grad
+#   dA_cumsum_tail: B1 -> B0   chunk-tail dA_cumsum grad
+#   dh_last       : (region input) cotangent of final_state (zero by default)
+# All fp32 (recurrent-state precision; resolved in _buffer_dtype). The final input
+# grads (dx/dB/dlog_decay/ddt/dh0/dC/dz/dD) are this brick's grad outputs.
+_MAMBA3_CHUNKED_BWD_HANDOFF_SUFFIXES = (
+    "dchunk_states",
+    "dstates",
+    "dinp_diag",
+    "dA_cumsum_y",
+    "dA_cumsum_tail",
+)
 
 
 def _emit_mamba3_model_brick_surfaces(
@@ -2042,9 +2061,112 @@ def _emit_mamba3_chunked_model_brick_surfaces(
             backward="owner_output",
         )
     )
+    # When the flag is ON the forward surfaces use ``owner_output`` (no AOT
+    # autograd derivation), so the 3 backward segments are emitted EXPLICITLY here
+    # with the cross-segment grad-handoff wiring (B2 -> B1 -> B0), the analytic
+    # transpose of the forward F2 -> F1 -> F0. This replaces the serial
+    # ``mamba3_mimo_bwd`` (which is not emitted when the flag is ON).
+    _emit_mamba3_chunked_bwd_model_brick_surfaces(brick, context)
     return _PathCBrickSurfaceLoweringResult(delta_output=delta)
 
 
+def _emit_mamba3_chunked_bwd_model_brick_surfaces(
+    brick: _ResolvedPathCModelBrick,
+    context: _PathCBrickSurfaceLoweringContext,
+) -> None:
+    """Emit the chunked B2/B1/B0 SSD-core BACKWARD as 3 region surfaces.
+
+    The analytic transpose of the forward F2/F1/F0 (design §2/§7 Stage 3), wired by
+    the per-brick caller-owned grad-handoff buffers (``dchunk_states`` B2->B1,
+    ``dstates`` B1->B0, ``dinp_diag``/``dA_cumsum_y`` B2->B0, ``dA_cumsum_tail``
+    B1->B0). REUSES the forward-materialized boundary states
+    (``cb / dA_cumsum / prev_states / summary_states``) instead of the 8x
+    checkpoint-replay (design §3/§6).
+
+    Inputs are the brick's delta cotangent (``{name}_delta_grad``), the forward
+    cache (the same projected SSD tensors + handoff buffers the forward staged),
+    and ``{name}_dh_last`` (final-state cotangent, zero by default). Outputs are the
+    brick's grad outputs for the SSD inputs (``{name}_x_grad`` / ``B`` / ``C`` /
+    ``dt`` / ``A`` via dlog_decay+ddt / ``z`` / ``D`` / ``h0``).
+
+    Each ``_bwd`` op-node compiles to its proven ``build_*_bwd_metal`` grid kernel
+    via the compile-site delegation interpose. RULE #1: on compile/parity failure
+    the builder RAISES (no silent serial fallback).
+    """
+    name = brick.name
+    proj = {
+        suffix: f"{name}_{suffix}"
+        for suffix in _MAMBA3_CHUNKED_PROJECTED_SSD_INPUT_SUFFIXES
+    }
+    hand = {
+        suffix: f"{name}_{suffix}"
+        for suffix in _MAMBA3_CHUNKED_HANDOFF_SUFFIXES
+    }
+    bwd = {
+        suffix: f"{name}_{suffix}"
+        for suffix in _MAMBA3_CHUNKED_BWD_HANDOFF_SUFFIXES
+    }
+    h0 = f"{name}_h0"
+    skip_d = f"{name}_D"
+    delta_grad = f"{name}_delta_grad"
+    dh_last = f"{name}_dh_last"
+    y = f"{name}_y"
+    z = f"{name}_z"
+
+    # B2 — output/Y transpose. Reads delta cotangent + forward cache; writes the
+    # output-side grads (dC/dz/dD final; dchunk_states/dinp_diag/dA_cumsum_y handoff;
+    # dx D-skip partial accumulated by B0).
+    context.surfaces.append(
+        FusionKernelSurface.path_c(
+            name=f"{name}_chunk_scan_combine_bwd",
+            op_name="mamba3_chunk_scan_combine_bwd",
+            inputs=(
+                delta_grad, hand["cb"], proj["x"], z, proj["dt"],
+                hand["dA_cumsum"], proj["C"], proj["B"], hand["prev_states"],
+                skip_d, y,
+            ),
+            outputs=(
+                f"{name}_C_grad", f"{name}_x_grad", f"{name}_z_grad",
+                bwd["dchunk_states"], bwd["dinp_diag"], bwd["dA_cumsum_y"],
+                f"{name}_D_grad",
+            ),
+            backward="owner_output",
+        )
+    )
+    # B1 — REVERSE inter-chunk combiner. Reads dchunk_states + prev_states + the
+    # final-state cotangent; writes dstates/dA_cumsum_tail handoff + dh0 (final).
+    context.surfaces.append(
+        FusionKernelSurface.path_c(
+            name=f"{name}_inter_chunk_recur_bwd",
+            op_name="mamba3_inter_chunk_recur_bwd",
+            inputs=(
+                bwd["dchunk_states"], hand["dA_cumsum"], dh_last,
+                hand["prev_states"],
+            ),
+            outputs=(bwd["dstates"], f"{name}_h0_grad", bwd["dA_cumsum_tail"]),
+            backward="owner_output",
+        )
+    )
+    # B0 — precompute transpose. Folds dstates + dinp_diag + the dA_cumsum grads
+    # into the final input grads dx (accumulated)/dB/dlog_decay/ddt.
+    context.surfaces.append(
+        FusionKernelSurface.path_c(
+            name=f"{name}_chunk_precompute_bwd",
+            op_name="mamba3_chunk_precompute_bwd",
+            inputs=(
+                bwd["dstates"], bwd["dinp_diag"], bwd["dA_cumsum_y"],
+                bwd["dA_cumsum_tail"], hand["dA_cumsum"], proj["x"], proj["B"],
+                proj["dt"], proj["A"],
+            ),
+            outputs=(
+                f"{name}_x_grad", f"{name}_B_grad", f"{name}_dt_grad",
+                f"{name}_A_grad",
+            ),
+            backward="owner_output",
+        )
+    )
+
+
 def _emit_m2rnn_model_brick_surfaces(
     brick: _ResolvedPathCModelBrick,
     context: _PathCBrickSurfaceLoweringContext,

cppmega_mlx/runtime/path_c_fusion_schedules.py

@@ -265,6 +265,20 @@
         "mamba3_final_state",
     }
 )
+# Backward grad-handoff buffers (B2 -> B1 -> B0; Stage 3). Force-spilled to named
+# ABI scratch (mirror of the forward set) so they land as stable named params
+# across the producer/consumer backward segments. All fp32 (resolved in
+# _buffer_dtype). Flag-gated (default OFF) — never appear in the serial backward.
+DESCRIPTOR_MAMBA3_CHUNKED_BWD_HANDOFF_ABI_BUFFERS = frozenset(
+    {
+        "mamba3_dh_last",
+        "mamba3_dchunk_states",
+        "mamba3_dstates",
+        "mamba3_dinp_diag",
+        "mamba3_dA_cumsum_y",
+        "mamba3_dA_cumsum_tail",
+    }
+)
 DESCRIPTOR_ROW_PHASED_BWD_SCRATCH_ABI_CANONICALS = frozenset(
     {
         "hidden",
@@ -725,13 +739,18 @@ class _ScheduleNodeView:
 def _path_c_schedule_node_execution_phase(node: Any) -> str:
     backward = str(getattr(node, "backward", ""))
     op_name = str(getattr(node, "op_name", ""))
-    # The chunked F0/F1/F2 SSD-core ops are FORWARD producers that carry NO
-    # synthesized AOT backward (they own their output gradient: backward=
-    # "owner_output"), but they must run in the FORWARD execution stage so the
-    # compile-site delegation interpose substitutes their grid kernel during the
-    # forward pass (RULE #1: the chunked forward is never grouped as a backward
-    # fragment). Backward stays serial/unchanged.
-    if op_name in _MAMBA3_CHUNKED_GRID_DELEGATION_OPS:
+    # The chunked F0/F1/F2 SSD-core FORWARD ops carry NO synthesized AOT backward
+    # (they own their output gradient: backward="owner_output"), but they must run
+    # in the FORWARD execution stage so the compile-site delegation interpose
+    # substitutes their grid kernel during the forward pass (RULE #1: the chunked
+    # forward is never grouped as a backward fragment). The chunked B0/B1/B2
+    # BACKWARD ops (``*_bwd``) are explicitly EXCLUDED here so they classify as
+    # backward below (they too are in _MAMBA3_CHUNKED_GRID_DELEGATION_OPS for the
+    # shared interpose, but they belong to the backward execution stage).
+    if (
+        op_name in _MAMBA3_CHUNKED_GRID_DELEGATION_OPS
+        and not op_name.endswith("_bwd")
+    ):
         return "forward"
     if backward == "owner_output" or op_name.endswith("_bwd"):
         return "backward"
@@ -4663,6 +4682,8 @@ def _is_row_phased_bwd_scratch_abi_buffer(buffer_name: str) -> bool:
         return True
     if name in DESCRIPTOR_MAMBA3_CHUNKED_FWD_HANDOFF_ABI_BUFFERS:
         return True
+    if name in DESCRIPTOR_MAMBA3_CHUNKED_BWD_HANDOFF_ABI_BUFFERS:
+        return True
     if not name.endswith("_grad"):
         return False
     canonical = _canonical_buffer_name(name)
@@ -14375,6 +14396,14 @@ def _buffer_dtype(
         "mamba3_summary_states",
         "mamba3_prev_states",
         "mamba3_final_state",
+        # Mamba3 chunked-scan B2->B1->B0 backward grad-handoff buffers (Stage 3):
+        # fp32 for recurrent grad-state precision.
+        "mamba3_dh_last",
+        "mamba3_dchunk_states",
+        "mamba3_dstates",
+        "mamba3_dinp_diag",
+        "mamba3_dA_cumsum_y",
+        "mamba3_dA_cumsum_tail",
         "m2rnn_h_state",
         "m2rnn_h_checkpoint",
         "m2rnn_h0_grad",
@@ -14502,6 +14531,13 @@ def _is_mamba3_state_like_buffer(buffer_name: str, canonical_name: str) -> bool:
     "summary_states",
     "prev_states",
     "final_state",
+    # backward grad-handoff (B2 -> B1 -> B0; Stage 3)
+    "dh_last",
+    "dchunk_states",
+    "dstates",
+    "dinp_diag",
+    "dA_cumsum_y",
+    "dA_cumsum_tail",
 )
 
 
@@ -14594,6 +14630,13 @@ def _mamba3_chunked_fwd_region_buffer_shape(
         "summary_states": (batch, nchunks, nheads, headdim, dstate),
         "prev_states": (batch, nchunks, nheads, headdim, dstate),
         "final_state": (batch, nheads, headdim, dstate),
+        # BACKWARD grad-handoff buffers (B2 -> B1 -> B0; Stage 3, fp32)
+        "dh_last": (batch, nheads, headdim, dstate),
+        "dchunk_states": (batch, nchunks, nheads, headdim, dstate),
+        "dstates": (batch, nchunks, nheads, headdim, dstate),
+        "dinp_diag": (batch, seqlen, nheads, headdim, dstate),
+        "dA_cumsum_y": (batch, nheads, nchunks, chunk),
+        "dA_cumsum_tail": (batch, nheads, nchunks, chunk),
     }
     return shapes[suffix]
 

tests/test_mamba3_chained_forward_f0f1f2.py

@@ -599,8 +599,13 @@ def test_region_flip_on_full_scale_region_parity(monkeypatch, capsys):
     assert env.mamba_num_heads == H
 
     region = _build_mamba_direct_chain_region(env)
+    # Forward-only chunked nodes (Stage 3 also emits 3 chunked _bwd nodes when the
+    # flag is ON; this forward-parity check exercises just F0/F1/F2).
     chunked_nodes = [
-        n for n in region.nodes if n.op_name in _MAMBA3_CHUNKED_GRID_DELEGATION_OPS
+        n
+        for n in region.nodes
+        if n.op_name in _MAMBA3_CHUNKED_GRID_DELEGATION_OPS
+        and not n.op_name.endswith("_bwd")
     ]
     assert len(chunked_nodes) == 3
 

tests/test_mamba3_chunked_backward_b0b1b2.py

@@ -170,6 +170,93 @@ def test_backward_interpose_on_emits_real_grid_kernel(op_name, monkeypatch):
 # --------------------------------------------------------------------------- #
 
 
+# --------------------------------------------------------------------------- #
+# Verification 1b: the LIVE REGION-BUILD backward 1->3 surface flip.            #
+# --------------------------------------------------------------------------- #
+
+
+def _build_mamba_direct_chain_region(env):
+    from cppmega_mlx.runtime.path_c_fusion import (
+        PathCModelBrick,
+        build_path_c_model_region_from_bricks,
+    )
+
+    bricks = (PathCModelBrick(name="mamba3_scan", kind="mamba3", route_symbol="M"),)
+    return build_path_c_model_region_from_bricks(
+        region_name="mamba3_direct_chain", bricks=bricks, shape_env=env
+    )
+
+
+def test_region_flip_off_no_chunked_backward_surfaces(monkeypatch):
+    """Flag OFF (default): the direct-chain region has NO chunked _bwd surfaces —
+    the serial mamba3_mimo (with aot_autograd) is unchanged (merge-safe)."""
+    monkeypatch.delenv("CPPMEGA_PATH_C_MAMBA3_CHUNKED_SCAN", raising=False)
+    region = _build_mamba_direct_chain_region(_env())
+    ops = [n.op_name for n in region.nodes]
+    assert "mamba3_mimo" in ops
+    assert not any(o in _BWD_OPS for o in ops), ops
+
+
+def test_region_flip_on_emits_three_chunked_backward_surfaces(monkeypatch):
+    """Flag ON: the region emits the 3 chunked _bwd segments (B2->B1->B0), each
+    classified as a BACKWARD-phase node isolated into its own backward stage, wired
+    by the per-brick grad-handoff buffers."""
+    monkeypatch.setenv("CPPMEGA_PATH_C_MAMBA3_CHUNKED_SCAN", "1")
+    from cppmega_mlx.runtime.path_c_fusion_schedules import (
+        _path_c_schedule_node_execution_phase,
+        plan_path_c_descriptor_stage_groups,
+    )
+
+    region = _build_mamba_direct_chain_region(_env())
+    bwd_nodes = {n.op_name: n for n in region.nodes if n.op_name in _BWD_OPS}
+    assert set(bwd_nodes) == set(_BWD_OPS), list(bwd_nodes)
+    # all 3 classify as backward phase
+    for n in region.nodes:
+        if n.op_name in _BWD_OPS:
+            assert _path_c_schedule_node_execution_phase(n) == "backward", n.op_name
+
+    # each _bwd op is isolated into its own backward stage group
+    op_by_node = {n.name: n.op_name for n in region.nodes}
+    bwd_stage_ops = []
+    for g in plan_path_c_descriptor_stage_groups(region):
+        ops = [op_by_node.get(x) for x in g.active_node_names]
+        for op in ops:
+            if op in _BWD_OPS:
+                assert g.execution_stage == "backward", g.execution_stage
+                assert len(g.active_node_names) == 1, list(g.active_node_names)
+                bwd_stage_ops.append(op)
+    assert bwd_stage_ops == list(_BWD_OPS), bwd_stage_ops
+
+    # grad-handoff wiring: B2 -> B1 -> B0 (the transpose of F2 -> F1 -> F0)
+    b2 = bwd_nodes["mamba3_chunk_scan_combine_bwd"]
+    b1 = bwd_nodes["mamba3_inter_chunk_recur_bwd"]
+    b0 = bwd_nodes["mamba3_chunk_precompute_bwd"]
+    assert "mamba3_scan_dchunk_states" in b2.outputs
+    assert "mamba3_scan_dchunk_states" in b1.inputs
+    assert "mamba3_scan_dstates" in b1.outputs
+    assert "mamba3_scan_dstates" in b0.inputs
+    assert "mamba3_scan_dinp_diag" in b2.outputs and "mamba3_scan_dinp_diag" in b0.inputs
+    assert "mamba3_scan_dA_cumsum_tail" in b1.outputs and "mamba3_scan_dA_cumsum_tail" in b0.inputs
+    # B2/B1/B0 REUSE the forward-materialized boundary states (no replay)
+    assert "mamba3_scan_prev_states" in b2.inputs
+    assert "mamba3_scan_prev_states" in b1.inputs
+    assert "mamba3_scan_cb" in b2.inputs
+
+
+def test_backward_handoff_buffers_fp32_and_force_spilled():
+    """The 6 backward grad-handoff buffers resolve fp32 and are in the force-spill
+    ABI set (mirror of the forward handoff registration)."""
+    from cppmega_mlx.runtime.path_c_fusion_schedules import (
+        DESCRIPTOR_MAMBA3_CHUNKED_BWD_HANDOFF_ABI_BUFFERS,
+        _buffer_dtype,
+    )
+
+    for nm in DESCRIPTOR_MAMBA3_CHUNKED_BWD_HANDOFF_ABI_BUFFERS:
+        assert _buffer_dtype(nm) == "float32", nm
+    assert "mamba3_dchunk_states" in DESCRIPTOR_MAMBA3_CHUNKED_BWD_HANDOFF_ABI_BUFFERS
+    assert "mamba3_dstates" in DESCRIPTOR_MAMBA3_CHUNKED_BWD_HANDOFF_ABI_BUFFERS
+
+
 @pytest.mark.skipif(
     not _torch_mps_available(), reason="requires torch + Metal (mps) backend"
 )