fix(mamba3-pathc): flag-ON direct-chain emits FULL backward (non-mamba bwd + chunked B2/B1/B0), seeds mamba delta cotangent

ROOT CAUSE: the direct-chain planner (plan_path_c_direct_fusion_chain_for_region)
skipped build_path_c_aot_autograd_region whenever ANY _bwd node was already in the
region. Flag ON appends the chunked mamba B2/B1/B0 _bwd surfaces into the forward
node list, so that guard suppressed the AOT-backward derivation for ALL OTHER
bricks -> the chain had only the 3 chunked mamba bwd (interleaved mid-forward) and
nothing seeded the mamba delta_grad cotangent (zero mamba grads).

FIX (shared by both flag states):
 - build_path_c_aot_autograd_region: drop explicit owner_output _bwd nodes from
   the forward list; _aot_backward_surfaces_for walks reversed(forward_nodes) and
   emits each chunked-mamba brick's pre-built B2/B1/B0 at that reverse position, so
   residual_rmsnorm_bwd (which produces {brick}_delta_grad) runs BEFORE the mamba
   backward and seeds its cotangent.
 - planner guard now short-circuits only on a SYNTHESIZED (non-owner_output) _bwd.

VERIFIED: flag-OFF chain unchanged (14 segs, byte-identical). flag-ON chain now 18
segs (9 fwd + 9 bwd) in correct reverse order; brick_10_M_delta_grad nonzero and
read by B2; live route grads 2 -> 18 (m2rnn+attention+suffix now flow). New tests:
full-backward-chain OFF/ON ordering + delta_grad seam (13 passed).

REMAINING (separate feature, NOT this change): 163-grad parity still blocked by the
absence of an eager mamba projection fwd/bwd bridge -- the chunked projected inputs
(x/B/C/A/dt/z/h0) are zero-seeded in the pre-step owner (no in_proj/conv/norm
forward), so chunked grads are zero AND are w.r.t. SSD-core inputs, not the
mamba3_*_weight model params the grad-tree aliases expect. Documented in
scratch/MAMBA3-RUNTIME-WIRING-RESULTS.md. RULE #1: did NOT alias mismatched grads.

Flag CPPMEGA_PATH_C_MAMBA3_CHUNKED_SCAN default OFF (merge-safe).

Author: apstenku123

cppmega_mlx/runtime/path_c_fusion.py

@@ -2475,27 +2475,127 @@ def _aot_backward_surface_for_node(node: FusionNode) -> FusionKernelSurface:
     )
 
 
-def _aot_backward_surfaces_for(region: PathCFusionRegion) -> tuple[FusionKernelSurface, ...]:
-    return tuple(
-        _aot_backward_surface_for_node(node)
-        for node in reversed(region.nodes)
-        if node.backward == "aot_autograd"
+# The chunked mamba3 SSD-core ops emit their OWN backward surfaces (owner_output)
+# rather than relying on the synthesized AOT ``_bwd`` derivation. The forward F0/F1/
+# F2 ops are forward-phase owner_output nodes; their explicit B2/B1/B0 ``_bwd`` ops
+# (in forward-list order B2,B1,B0 -- already the correct reverse-of-forward order)
+# are appended into the brick's surface list by
+# ``_emit_mamba3_chunked_bwd_model_brick_surfaces``.
+_MAMBA3_CHUNKED_FORWARD_OP_NAMES = (
+    "mamba3_chunk_precompute",
+    "mamba3_inter_chunk_recur",
+    "mamba3_chunk_scan_combine",
+)
+
+
+def _is_mamba3_chunked_forward_node(node: FusionNode) -> bool:
+    return (
+        str(node.op_name) in _MAMBA3_CHUNKED_FORWARD_OP_NAMES
+        and str(node.backward) == "owner_output"
+    )
+
+
+def _is_owner_output_backward_node(node: FusionNode) -> bool:
+    """True for an EXPLICIT backward surface already present in a forward region.
+
+    These are the chunked mamba3 B2/B1/B0 ``_bwd`` ops emitted alongside the
+    forward F0/F1/F2 surfaces (op_name ends ``_bwd`` and backward=owner_output).
+    They are removed from the forward-node list and re-emitted, in correct
+    reverse-of-forward position, by ``_aot_backward_surfaces_for``.
+    """
+
+    return (
+        str(node.op_name).endswith("_bwd")
+        and str(node.backward) == "owner_output"
     )
 
 
+def _explicit_backward_nodes_for_forward(
+    forward_node: FusionNode,
+    region: PathCFusionRegion,
+) -> tuple[FusionNode, ...]:
+    """Return the explicit ``_bwd`` nodes belonging to a chunked-mamba brick.
+
+    The chunked mamba3 brick named ``{brick}`` lowers F0/F1/F2 with names
+    ``{brick}_chunk_precompute`` / ``_inter_chunk_recur`` / ``_chunk_scan_combine``
+    and the matching B2/B1/B0 with the SAME stems + ``_bwd``. The three backward
+    surfaces are attached to the F0 (first) forward node so the reverse derivation
+    emits them once per brick, in their forward-list (B2,B1,B0 = reverse) order,
+    at the reverse position of the brick's forward span.
+    """
+
+    name = str(forward_node.name)
+    if not name.endswith("_chunk_precompute"):
+        return ()
+    brick = name[: -len("_chunk_precompute")]
+    bwd_stems = (
+        f"{brick}_chunk_scan_combine_bwd",
+        f"{brick}_inter_chunk_recur_bwd",
+        f"{brick}_chunk_precompute_bwd",
+    )
+    by_name = {str(node.name): node for node in region.nodes}
+    selected = []
+    for stem in bwd_stems:
+        node = by_name.get(stem)
+        if node is None:
+            raise ValueError(
+                "chunked mamba3 backward surface missing for forward brick "
+                f"{brick!r}: expected {stem!r} in region {region.name!r}"
+            )
+        selected.append(node)
+    return tuple(selected)
+
+
+def _aot_backward_surfaces_for(region: PathCFusionRegion) -> tuple[FusionKernelSurface, ...]:
+    """Return backward surfaces in reverse-of-forward execution order.
+
+    For each FORWARD node (reversed): an ``aot_autograd`` node derives a
+    synthesized ``_bwd`` surface; a chunked-mamba forward (owner_output) node
+    emits its pre-built explicit B2/B1/B0 ``_bwd`` surfaces at that reverse
+    position (so the model graph's later-in-reverse residual_rmsnorm_bwd, which
+    produces the brick ``delta_grad`` cotangent, runs BEFORE the mamba backward).
+    Explicit ``_bwd`` nodes already in the region are NOT re-derived as forward.
+    """
+
+    surfaces: list[FusionKernelSurface] = []
+    for node in reversed(region.nodes):
+        if _is_owner_output_backward_node(node):
+            # Explicit backward surfaces are emitted via their forward node below.
+            continue
+        if node.backward == "aot_autograd":
+            surfaces.append(_aot_backward_surface_for_node(node))
+        elif _is_mamba3_chunked_forward_node(node):
+            surfaces.extend(
+                _surface_from_node(bwd_node)
+                for bwd_node in _explicit_backward_nodes_for_forward(node, region)
+            )
+    return tuple(surfaces)
+
+
 def build_path_c_aot_autograd_region(
     region: PathCFusionRegion,
     *,
     region_name: str | None = None,
 ) -> PathCFusionRegion:
-    """Return ``region`` plus symbolic AOT backward graph surfaces."""
+    """Return ``region`` plus symbolic AOT backward graph surfaces.
+
+    The forward node list keeps every FORWARD-phase surface (including the chunked
+    mamba3 F0/F1/F2 owner_output forward ops) but drops the explicit ``_bwd``
+    surfaces that the chunked brick appended alongside its forward; those are
+    re-emitted, in correct reverse position, by ``_aot_backward_surfaces_for``.
+    """
 
     if not isinstance(region, PathCFusionRegion):
         raise TypeError("region must be PathCFusionRegion")
+    forward_surfaces = tuple(
+        _surface_from_node(node)
+        for node in region.nodes
+        if not _is_owner_output_backward_node(node)
+    )
     return build_path_c_fusion_region(
         region_name=region_name or region.name,
         surfaces=(
-            *(_surface_from_node(node) for node in region.nodes),
+            *forward_surfaces,
             *_aot_backward_surfaces_for(region),
         ),
         z3_sync=region.z3_sync,

cppmega_mlx/runtime/path_c_fusion_schedules.py

@@ -16041,10 +16041,21 @@ def plan_path_c_direct_fusion_chain_for_region(
         _path_c_default_target() == "metal"
         and forward_max_segment_nodes is not None
     )
+    # Build the symbolic backward graph when requested. A region whose nodes
+    # ALREADY include a synthesized (aot-derived) backward op is treated as
+    # already-built and left as-is. The chunked mamba3 brick appends its OWN
+    # explicit B2/B1/B0 ``_bwd`` surfaces (backward="owner_output") into the
+    # FORWARD node list; those must NOT short-circuit the backward build --
+    # ``build_path_c_aot_autograd_region`` relocates them to their correct
+    # reverse position AND derives the remaining non-mamba brick backward ops.
+    already_has_synthesized_backward = any(
+        node.op_name.endswith("_bwd")
+        and str(getattr(node, "backward", "")) != "owner_output"
+        for node in region.nodes
+    )
     working_region = (
         build_path_c_aot_autograd_region(region)
-        if include_backward
-        and not any(node.op_name.endswith("_bwd") for node in region.nodes)
+        if include_backward and not already_has_synthesized_backward
         else region
     )
     nodes = tuple(working_region.nodes)

scratch/MAMBA3-RUNTIME-WIRING-RESULTS.md

@@ -189,3 +189,94 @@ the model gradient tree. Both are chain/grad-tree assembly — NOT a dtype/ABI b
     SAME 3 pre-existing failures (verified identical via `git stash` on clean base:
     test_fp8_..._blocks_missing_sparse_mla_producer + the 2 direct-chain
     value_and_grad bridge tests; unrelated to this change).
+
+================================================================================
+## UPDATE (commit on `mamba3-grad-tree`) — NON-MAMBA BACKWARD SEGMENTS NOW EMIT
+## under flag ON; the mamba delta cotangent IS seeded. The PRIOR "NEXT GAP"
+## (missing non-mamba bwd segs) is RESOLVED. A DEEPER gap blocks 163-grad parity.
+================================================================================
+
+### Root cause of the missing non-mamba backward segments (CORRECTS the premise above)
+The flag-ON `build_path_c_model_regions_from_model(include_backward=True)` region
+ALREADY emits ALL 9 backward nodes in correct reverse order (verified). The drop
+was in the DIRECT-CHAIN PLANNER: `plan_path_c_direct_fusion_chain_for_region`
+(path_c_fusion_schedules.py ~:16044) guarded the AOT-backward build with
+`not any(node.op_name.endswith("_bwd") for node in region.nodes)`. When the flag is
+ON the chunked brick appends its 3 explicit B2/B1/B0 `_bwd` surfaces INTO the
+forward node list, so that guard saw a `_bwd` node and SKIPPED
+`build_path_c_aot_autograd_region` ENTIRELY — so the OTHER bricks' backward
+(sparse_mla/attention/residual x2/m2rnn/entry) were never derived. The chain had
+only the 3 chunked mamba bwd, interleaved mid-forward.
+
+### FIX (gated nowhere — both flag states share the same correct code path)
+1. `build_path_c_aot_autograd_region` (path_c_fusion.py) now: (a) drops explicit
+   owner_output `_bwd` nodes from the forward list, (b) `_aot_backward_surfaces_for`
+   walks `reversed(forward_nodes)` and, for each chunked-mamba FORWARD node, emits
+   that brick's pre-built B2/B1/B0 `_bwd` surfaces AT THAT REVERSE POSITION (so the
+   residual_rmsnorm_bwd that produces `{brick}_delta_grad` runs BEFORE the mamba
+   backward). aot_autograd nodes derive their `_bwd` as before.
+2. The planner guard now only short-circuits on a SYNTHESIZED (non-owner_output)
+   `_bwd` node, so the chunked explicit `_bwd` no longer suppresses the build.
+
+### VERIFIED (probe + live route)
+- Flag-OFF chain UNCHANGED: 14 segments (7 fwd + 7 bwd), byte-identical order.
+- Flag-ON chain NOW: 18 segments (9 fwd + 9 bwd). Backward order:
+  sparse_mla_bwd, attention_qkv_bwd, residual_rmsnorm_bwd(R), m2rnn_bwd,
+  residual_rmsnorm_bwd(M -> emits brick_10_M_delta_grad), B2, B1, B0,
+  entry_rmsnorm_bwd. seg13 residual_rmsnorm_bwd OUT=brick_10_M_delta_grad;
+  seg14 B2 IN=brick_10_M_delta_grad (cotangent SEEDED).
+- Live flag-ON route end-to-end: loss FINITE (5.65–5.78), grads jumped 2 -> 18
+  (m2rnn layer-11 + attention layer-12 + suffix param grads now ALL flow because
+  the non-mamba backward runs). `brick_10_M_delta_grad` absmax=5.64e-3 (NONZERO).
+- New tests (test_mamba3_chunked_backward_b0b1b2.py): full-backward-chain OFF +
+  ON ordering + delta_grad seam — 13 passed (11 prior + 2 new).
+
+### THE DEEPER REMAINING GAP — chunked path has NO eager projection fwd/bwd bridge
+163-grad PARITY is STILL NOT achieved, but the cause is NOT the backward chain.
+The chunked region only covers the SSD CORE. Its projected inputs
+(`{brick}_x/B/C/A/dt/z/h0`) are caller-owned and the pre-step owner ALLOCATES THEM
+AS ZEROS (m04 `make_path_c_direct_chain_pre_step_runtime_owner` ~:8699) — there is
+NO eager `in_proj/conv/norm` forward that computes them from the residual hidden,
+and NO eager projection backward that folds the SSD-core input-grads back into the
+model params. Consequences, both verified at flag-ON:
+  (a) The chunked FORWARD runs on ZERO x/B/C/A/dt, so the cached boundary states
+      (cb/dA_cumsum/prev_states) are degenerate. The B2/B1/B0 backward multiplies
+      the (now NONZERO) `delta_grad` by zero forward activations -> EVERY chunked
+      grad output is ZERO: brick_10_M {x,B,C,A,dt,z,D,h0}_grad all absmax=0.
+  (b) Even if nonzero, those are grads w.r.t. the PROJECTED SSD tensors, NOT the
+      model params. The serial `mamba3_mimo_bwd` emits `{brick}_mamba3_in_proj_
+      weight_grad`/`_conv_weight_grad`/... directly (it backprops proj+conv+norm
+      internally). The grad-tree alias (`path_c_parameter_gradient_aliases`,
+      hybrid_lm.py:991) maps the mamba PARAMS to those `mamba3_*_weight_grad`
+      buffers — which the chunked chain never produces — so the coverage gate
+      (`_path_c_direct_chain_full_gradient_coverage_payload`) is INCOMPLETE for
+      layer-10 and the full-grad install path is bypassed (18 grads, not 163).
+  Only `{brick}_D_grad` would map 1:1 to a model param (mamba3_D); the projection/
+  conv/norm params have no chunked producer at all.
+
+RULE #1: I did NOT fake parity by aliasing the mismatched SSD-input grads onto the
+projection-param grad names (that would be a silent wrong-gradient fallback). The
+honest remaining work to reach 163-grad parity is a SEPARATE feature:
+  1. Eager (MLX-differentiable) mamba in_proj/conv/B-C-norm forward that POPULATES
+     `{brick}_x/B/C/A/dt/z/h0` from the residual hidden + the real block weights
+     (replacing the zero-seed), staged ahead of the chunked region.
+  2. Eager projection BACKWARD: take the chunked B0/B1/B2 SSD-input grads
+     (`x_grad/B_grad/.../z_grad/h0_grad`) and VJP them through that eager
+     projection to produce `{brick}_mamba3_*_weight_grad` model-param grads, then
+     register those under the existing grad-tree aliases so coverage completes.
+  This is analogous to the existing `path_c_prefix_gradient_tree_from_hidden_
+  cotangent` bridge but for the mamba projection sub-graph. It is NOT a chain-
+  assembly or grad-name-mapping change; it is a missing differentiable sub-step.
+
+### wall-time (Metal, local_gb10_quarter smoke, seq=128, batch=1)
+  flag OFF (serial): 163 grads, loss 5.7134, fwd+bwd 6.106 s.
+  flag ON (chunked): 18 grads (partial), loss ~5.65–5.78 finite, fwd+bwd ~26 s
+    (the verify harness builds+compiles 18 segments incl. the now-present 9 bwd;
+    the chunked GPU kernels themselves remain ~3x faster — the wall-time here is
+    dominated by first-call compile of the 4 extra backward segments).
+
+### Merge-safety re-confirmed with THIS change (flag default OFF)
+  - test_path_c_fusion_ir.py + test_path_c_autosplit_metal_parity.py: 124 passed.
+  - test_mamba3_chunked_backward_b0b1b2.py: 13 passed (incl. 2 new chain tests).
+  - test_m04_train_step.py -k direct_chain|path_c|chain: 82 passed, SAME 3
+    pre-existing failures (re-verified identical via `git stash` on clean base).

scratch/probe_chain_segs.py

@@ -0,0 +1,42 @@
+"""Probe the DIRECT CHAIN segments (planner output) for flag OFF vs ON."""
+from __future__ import annotations
+import os, sys
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), os.pardir, "scripts"))
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), os.pardir))
+
+import mlx.core as mx
+import numpy as np
+import m04_train_step as m04  # type: ignore
+from cppmega_mlx.recipes.model_factory import build_local_gb10_quarter_tiny_smoke_model
+from cppmega_mlx.runtime.path_c_fusion import path_c_mamba3_chunked_scan_enabled
+
+flag_on = path_c_mamba3_chunked_scan_enabled()
+mode = "ON" if flag_on else "OFF"
+print(f"=== flag {mode} ===")
+
+mx.random.seed(0)
+model = build_local_gb10_quarter_tiny_smoke_model(
+    hidden_size=64, num_attention_heads=1, mamba_expand=1, mamba_head_dim=64,
+    mamba_state_dim=16, mamba_groups=1, mamba_chunk_size=64,
+)
+profile_name = str(getattr(model, "path_c_profile_name", "HybridTinyLM"))
+prefix = m04._path_c_direct_chain_region_prefix(model, profile_name)
+seq = 128
+
+direct_chains = m04.plan_path_c_direct_fusion_chains_for_model(
+    model, region_prefix=prefix, include_backward=True, max_segment_nodes=1,
+    sequence_length=seq,
+)
+regions = m04.build_path_c_model_regions_from_model(
+    model, region_prefix=prefix, include_backward=False, sequence_length=seq,
+)
+sel_region = m04._select_path_c_model_route_region(regions)
+chain = m04._select_path_c_direct_chain_for_region(direct_chains, sel_region)
+print(f"chain status={getattr(chain,'status',None)} n_segments={len(chain.segments)}")
+for seg in chain.segments:
+    region = seg.region
+    ops = []
+    for n in getattr(region, "nodes", ()):
+        ops.append(n.op_name)
+    phase = str(getattr(seg, "execution_phase", "?"))
+    print(f"  seg[{seg.index}] status={seg.status} phase={phase} region={region.name} ops={ops}")

tests/test_mamba3_chunked_backward_b0b1b2.py

@@ -243,6 +243,87 @@ def test_region_flip_on_emits_three_chunked_backward_surfaces(monkeypatch):
     assert "mamba3_scan_cb" in b2.inputs
 
 
+def _build_mra_full_backward_region(env, *, include_backward):
+    """Build a 3-brick (M mamba3 / R m2rnn / A attention) region.
+
+    This mirrors the real model route span the direct-chain planner discovers so
+    the FULL reverse backward chain (non-mamba ``_bwd`` + chunked mamba B2/B1/B0)
+    can be asserted.
+    """
+    from cppmega_mlx.runtime.path_c_fusion import (
+        PathCModelBrick,
+        build_path_c_model_region_from_bricks,
+    )
+
+    bricks = (
+        PathCModelBrick(name="brick_M", kind="mamba3", route_symbol="M"),
+        PathCModelBrick(name="brick_R", kind="m2rnn", route_symbol="R"),
+        PathCModelBrick(name="brick_A", kind="attention", route_symbol="A"),
+    )
+    return build_path_c_model_region_from_bricks(
+        region_name="mra_chain",
+        bricks=bricks,
+        shape_env=env,
+        include_backward=include_backward,
+    )
+
+
+def test_full_backward_chain_off_emits_serial_mamba_bwd(monkeypatch):
+    """Flag OFF: the full (fwd+bwd) region emits the 7 serial backward ops in
+    reverse order, ending with mamba3_mimo_bwd then entry_rmsnorm_bwd. Merge-safe
+    baseline for the flag-ON assertion below."""
+    monkeypatch.delenv("CPPMEGA_PATH_C_MAMBA3_CHUNKED_SCAN", raising=False)
+    region = _build_mra_full_backward_region(_env(), include_backward=True)
+    bwd = [n.op_name for n in region.nodes if n.op_name.endswith("_bwd")]
+    assert bwd == [
+        "sparse_mla_fp8_apply_bwd",
+        "attention_qkv_projection_bwd",
+        "residual_rmsnorm_bwd",
+        "m2rnn_bwd",
+        "residual_rmsnorm_bwd",
+        "mamba3_mimo_bwd",
+        "entry_rmsnorm_bwd",
+    ], bwd
+
+
+def test_full_backward_chain_on_emits_all_segments_in_reverse_order(monkeypatch):
+    """Flag ON (the unlock fix): the full (fwd+bwd) region emits the NON-mamba
+    backward segments (sparse_mla/attention/residual x2/m2rnn/entry) IN ADDITION
+    TO the 3 chunked mamba ``_bwd`` (B2/B1/B0), all in correct reverse-of-forward
+    order. The residual_rmsnorm_bwd producing the mamba ``brick_M_delta_grad``
+    cotangent MUST run BEFORE the chunked mamba backward so the cotangent is
+    seeded (the root cause of the prior zero-mamba-grad gap)."""
+    monkeypatch.setenv("CPPMEGA_PATH_C_MAMBA3_CHUNKED_SCAN", "1")
+    region = _build_mra_full_backward_region(_env(), include_backward=True)
+    bwd = [n.op_name for n in region.nodes if n.op_name.endswith("_bwd")]
+    assert bwd == [
+        "sparse_mla_fp8_apply_bwd",
+        "attention_qkv_projection_bwd",
+        "residual_rmsnorm_bwd",
+        "m2rnn_bwd",
+        "residual_rmsnorm_bwd",
+        "mamba3_chunk_scan_combine_bwd",
+        "mamba3_inter_chunk_recur_bwd",
+        "mamba3_chunk_precompute_bwd",
+        "entry_rmsnorm_bwd",
+    ], bwd
+
+    # The mamba cotangent seam: a residual_rmsnorm_bwd must OUTPUT the brick's
+    # delta_grad and run strictly BEFORE the chunked mamba B2 reads it.
+    nodes = list(region.nodes)
+    delta_grad = "brick_M_delta_grad"
+    seed_idx = next(
+        i for i, n in enumerate(nodes)
+        if n.op_name == "residual_rmsnorm_bwd" and delta_grad in n.outputs
+    )
+    b2_idx = next(
+        i for i, n in enumerate(nodes)
+        if n.op_name == "mamba3_chunk_scan_combine_bwd"
+    )
+    assert seed_idx < b2_idx, (seed_idx, b2_idx)
+    assert delta_grad in nodes[b2_idx].inputs, nodes[b2_idx].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)."""