feat(path_c): device-resident gridded BACKWARD (B0/B1/B2) — CUDA pass_configs + delegation wire + RULE #1 no-numpy-fallback + gb10 probe

apstenku123 · apstenku123 · commit 6969a644edea · 2026-06-03T16:11:52.000+02:00
Make the MR region's backward device-resident gridded (mirrors the forward 77cb4a6+b961c6d), removing the abstract numpy host backward (the §15 79.6%@8l / 91.4%@28l wall) as a SILENT fallback. 1. CUDA pass_configs (mamba3_chunked_backward_core.py): thread pass_configs={tl.disable_tma_lower:True, tl.disable_warp_specialized:True} into each of build_chunk_scan_combine_bwd_metal / build_inter_chunk_recur_bwd_metal / build_chunk_precompute_bwd_metal ONLY on the resolved CUDA branch, mirroring compile_chunk_scan_fwd_metal exactly. The B0/B1/B2 prim BODIES have ZERO T.gemm / shared.dyn / make_swizzled_layout (grep-confirmed) so NO new CUDA sibling prim is needed. Metal default path is byte-identical (else branch unchanged). 2. Backward delegation: the existing interpose already routes the 3 bwd ops (_MAMBA3_CHUNKED_GRID_DELEGATION_OPS, descriptors, production_source, fragment_emitters, BWD handoff ABI buffers) via the direct-chain region path with the flag ON. Added register_real_backward_driver + make_real_bank_backward_driver (path_c_relax_train_step.py), symmetric to register_real_forward_driver, to drive the real gridded backward MR kernel (run_backward=1) over the bank ABI for the bank-SSA train_step path (device-resident, no host bounce). 3. RULE #1 (no silent numpy fallback): register_bank_drivers (path_c_relax_step_banks.py) now RAISES with WHERE+WHAT when the chunked flag is ON and the real gridded bank backward is NOT installed — instead of silently binding pathc.bank_bwd_* to the numpy _region_bwd_driver. set_real_bank_bwd_installed clears the guard once register_real_backward_driver runs (and stops the numpy bwd from clobbering the real binding). Explicit CPU self-check escape: CPPMEGA_PATH_C_ALLOW_NUMPY_BANK_BWD=1. Flag OFF (default) = byte-identical to before. 4. Probe scratch/probe_chunked_backward_cuda_gb10.py: prod shape (b=1,S=4096,c=64, G=8,H=112,P=64,N=64); seeds deterministic inputs + cotangent; builds the gridded forward cache (F0/F1, target=cuda) + serial y; runs B2->B1->B0 (pre-zeroed fp32 out, dx accumulated B2 D-skip then B0 inp-path); reference = MLX-free torch.autograd VJP of OUR exact serial recurrence (GOLD oracle); asserts all 8 model-facing grads max|abs|<1e-3 over ALL elements + NaN/inf guard; times B0/B1/B2 per-call; prints a machine-parseable RESULT json. RULE #1: any grad over gate or NaN RAISES. NO GPU on this Mac: edits + py_compile/AST + cross-module symbol checks only; gb10 compile/run is the Profile agent's single-run step.
diff --git a/cppmega_mlx/nn/_tilelang/mamba3_chunked_backward_core.py b/cppmega_mlx/nn/_tilelang/mamba3_chunked_backward_core.py
@@ -419,10 +419,31 @@ def build_chunk_scan_combine_bwd_metal(
     prim = chunk_scan_combine_bwd_metal_prim(
         batch, seqlen, chunk_size, ngroups, nheads, headdim, dstate, **kwargs
     )
+    resolved_target = _resolve_chunked_compile_target(target)
+    # CUDA (sm_121): mirror the forward F2 compile-site EXACTLY — thread the same
+    # two pass_configs so the CUDA codegen surface is identical to the validated
+    # forward. The B2 prim BODY has NO T.gemm / TMA-eligible copy (grep-confirmed:
+    # zero T.gemm / shared.dyn / make_swizzled_layout), so there is no tensormap
+    # descriptor to mis-align; disabling the TMA + warp-specialized lowering is a
+    # no-op-or-safer escape hatch that keeps the compile path byte-identical to
+    # the forward. The Metal branch is UNCHANGED (no pass_configs). RULE #1:
+    # explicit per-target codegen choice, never a silent fallback.
+    kind = str(getattr(getattr(resolved_target, "kind", None), "name", "")).lower()
+    if "cuda" in kind:
+        pass_configs = {
+            "tl.disable_tma_lower": True,
+            "tl.disable_warp_specialized": True,
+        }
+        return tilelang.compile(
+            prim,
+            out_idx=[11, 12, 13, 14, 15, 16, 17],
+            target=resolved_target,
+            pass_configs=pass_configs,
+        )
     return tilelang.compile(
         prim,
         out_idx=[11, 12, 13, 14, 15, 16, 17],
-        target=_resolve_chunked_compile_target(target),
+        target=resolved_target,
     )
 
 
@@ -596,10 +617,31 @@ def build_inter_chunk_recur_bwd_metal(
     prim = inter_chunk_recur_bwd_metal_prim(
         batch, seqlen, chunk_size, ngroups, nheads, headdim, dstate, **kwargs
     )
+    resolved_target = _resolve_chunked_compile_target(target)
+    # CUDA (sm_121): mirror the forward F2 compile-site EXACTLY — thread the same
+    # two pass_configs. The B1 prim BODY has NO T.gemm / TMA-eligible copy
+    # (grep-confirmed: zero T.gemm / shared.dyn / make_swizzled_layout — it does
+    # not even alloc_shared), so there is no tensormap descriptor to mis-align;
+    # the pass_configs are a no-op-or-safer escape hatch keeping the CUDA compile
+    # path byte-identical to the forward. The Metal branch is UNCHANGED (no
+    # pass_configs). RULE #1: explicit per-target codegen choice, never a silent
+    # fallback.
+    kind = str(getattr(getattr(resolved_target, "kind", None), "name", "")).lower()
+    if "cuda" in kind:
+        pass_configs = {
+            "tl.disable_tma_lower": True,
+            "tl.disable_warp_specialized": True,
+        }
+        return tilelang.compile(
+            prim,
+            out_idx=[4, 5, 6],
+            target=resolved_target,
+            pass_configs=pass_configs,
+        )
     return tilelang.compile(
         prim,
         out_idx=[4, 5, 6],
-        target=_resolve_chunked_compile_target(target),
+        target=resolved_target,
     )
 
 
@@ -852,8 +894,29 @@ def build_chunk_precompute_bwd_metal(
     prim = chunk_precompute_bwd_metal_prim(
         batch, seqlen, chunk_size, ngroups, nheads, headdim, dstate, **kwargs
     )
+    resolved_target = _resolve_chunked_compile_target(target)
+    # CUDA (sm_121): mirror the forward F2 compile-site EXACTLY — thread the same
+    # two pass_configs. The B0 prim BODY has NO T.gemm / TMA-eligible copy
+    # (grep-confirmed: zero T.gemm / shared.dyn / make_swizzled_layout; all
+    # alloc_shared are plain default scope), so there is no tensormap descriptor
+    # to mis-align; the pass_configs are a no-op-or-safer escape hatch keeping the
+    # CUDA compile path byte-identical to the forward. The Metal branch is
+    # UNCHANGED (no pass_configs). RULE #1: explicit per-target codegen choice,
+    # never a silent fallback.
+    kind = str(getattr(getattr(resolved_target, "kind", None), "name", "")).lower()
+    if "cuda" in kind:
+        pass_configs = {
+            "tl.disable_tma_lower": True,
+            "tl.disable_warp_specialized": True,
+        }
+        return tilelang.compile(
+            prim,
+            out_idx=[9, 10, 11, 12],
+            target=resolved_target,
+            pass_configs=pass_configs,
+        )
     return tilelang.compile(
         prim,
         out_idx=[9, 10, 11, 12],
-        target=_resolve_chunked_compile_target(target),
+        target=resolved_target,
     )
diff --git a/cppmega_mlx/runtime/path_c_relax_step_banks.py b/cppmega_mlx/runtime/path_c_relax_step_banks.py
@@ -84,8 +84,10 @@
     parse_physical_bank_shapes,
 )
 from cppmega_mlx.runtime.path_c_fusion import (
+    MAMBA3_CHUNKED_SCAN_ENV,
     build_path_c_aot_autograd_region,
     build_path_c_model_region_from_route_symbols,
+    path_c_mamba3_chunked_scan_enabled,
 )
 from cppmega_mlx.runtime.path_c_fusion_schedules import path_c_fusion_schedule_template
 from cppmega_mlx.recipes.model_factory import local_gb10_quarter_profile
@@ -295,12 +297,87 @@ def packed(actg_in, param, state_ckpt, paramg_in, actg_out, paramg_out):
     return packed
 
 
+# Env that EXPLICITLY acknowledges the abstract numpy bank backward driver is the
+# intended path for THIS bank-SSA build (the CPU self-check / non-gridded path).
+# Default OFF. When the gridded chunked-scan backward (B0/B1/B2) is enabled but the
+# real backward driver has NOT been installed, register_bank_drivers RAISES rather
+# than silently binding the numpy host backward (RULE #1: no silent fallback). Set
+# this truthy ONLY for the deliberate numpy self-check, never in production.
+ALLOW_NUMPY_BANK_BWD_ENV = "CPPMEGA_PATH_C_ALLOW_NUMPY_BANK_BWD"
+
+# Set truthy by register_real_backward_driver once the REAL gridded bank backward is
+# installed, so a subsequent register_bank_drivers (e.g. re-entry) does NOT clobber
+# the real bwd binding with the numpy one (and the RULE #1 guard knows the real path
+# is live). This is a process-local in-memory flag, NOT an env read.
+_REAL_BANK_BWD_INSTALLED = False
+
+
+def _numpy_bank_bwd_allowed() -> bool:
+    import os as _os
+
+    raw = _os.environ.get(ALLOW_NUMPY_BANK_BWD_ENV, "")
+    return raw.strip().lower() in {"1", "true", "yes", "on"}
+
+
+def set_real_bank_bwd_installed(value: bool) -> None:
+    """Mark whether the REAL gridded bank backward driver has been installed.
+
+    Called by ``register_real_backward_driver`` so a later ``register_bank_drivers``
+    does NOT clobber the real gridded ``pathc.bank_bwd_*`` binding with the abstract
+    numpy one, and so the RULE #1 guard treats the gridded backward as live (it does
+    NOT raise — the real path is installed, not the numpy fallback).
+    """
+    global _REAL_BANK_BWD_INSTALLED
+    _REAL_BANK_BWD_INSTALLED = bool(value)
+
+
+def real_bank_bwd_installed() -> bool:
+    """Return whether the REAL gridded bank backward driver is currently installed."""
+    return _REAL_BANK_BWD_INSTALLED
+
+
 def register_bank_drivers(numels: dict[str, int], n_layers: int) -> None:
+    # RULE #1 (NO SILENT FALLBACK): when the gridded chunked-scan backward is
+    # ENABLED, the bank-SSA train_step path's per-region `pathc.bank_bwd_i` is the
+    # ABSTRACT NUMPY host backward (`_region_bwd_driver`) — a host round-trip of the
+    # ~2028 MB region banks that does NOT route through the gridded B0/B1/B2 region
+    # surfaces. Binding it while the gridded backward is enabled would be a SILENT
+    # numpy-host fallback masquerading as the gridded path. So:
+    #   * If the real gridded bank backward is already installed
+    #     (`register_real_backward_driver` set `_REAL_BANK_BWD_INSTALLED`), keep it —
+    #     bind ONLY the forward here, never clobber the real bwd with numpy.
+    #   * Else if the flag is ON and numpy-bwd was NOT explicitly acknowledged
+    #     (`ALLOW_NUMPY_BANK_BWD_ENV`), RAISE with WHERE+WHAT — direct the caller to
+    #     the direct-chain region path (build_path_c_model_region_from_route_symbols
+    #     + compile_path_c_region, flag ON) OR register_real_backward_driver, which
+    #     actually execute the gridded B0/B1/B2.
+    #   * Else (flag OFF, or numpy explicitly acknowledged for the CPU self-check):
+    #     the numpy bank backward is the legitimate path — bind it.
+    chunked_on = path_c_mamba3_chunked_scan_enabled()
+    if chunked_on and not _REAL_BANK_BWD_INSTALLED and not _numpy_bank_bwd_allowed():
+        raise RuntimeError(
+            "register_bank_drivers (path_c_relax_step_banks): the gridded chunked "
+            "backward is ENABLED ("
+            f"{MAMBA3_CHUNKED_SCAN_ENV}=1) but this bank-SSA "
+            "build_train_step path would bind `pathc.bank_bwd_*` to the ABSTRACT "
+            "NUMPY host backward (_region_bwd_driver), which does NOT execute the "
+            "gridded B0/B1/B2 region surfaces. Binding it would be a SILENT "
+            "numpy-host backward fallback (RULE #1 forbidden). Drive the gridded "
+            "backward via EITHER the direct-chain region path "
+            "(build_path_c_model_region_from_route_symbols + compile_path_c_region "
+            "with the flag ON, which routes through the chunked region surfaces and "
+            "the delegation interpose) OR call register_real_backward_driver(...) to "
+            "install the real gridded bank backward. To DELIBERATELY run the numpy "
+            "self-check with the flag on, set "
+            f"{ALLOW_NUMPY_BANK_BWD_ENV}=1 (CPU reference only, never production)."
+        )
     for i in range(n_layers):
         tvm_ffi.register_global_func(
             f"pathc.bank_fwd_{i}", _region_fwd_driver(numels), override=True)
-        tvm_ffi.register_global_func(
-            f"pathc.bank_bwd_{i}", _region_bwd_driver(numels), override=True)
+        # Do NOT overwrite a real gridded bwd binding with the numpy one.
+        if not _REAL_BANK_BWD_INSTALLED:
+            tvm_ffi.register_global_func(
+                f"pathc.bank_bwd_{i}", _region_bwd_driver(numels), override=True)
 
 
 def build_bank_chain(numels: dict[str, int], n_layers: int) -> tvm.IRModule:
diff --git a/cppmega_mlx/runtime/path_c_relax_train_step.py b/cppmega_mlx/runtime/path_c_relax_train_step.py
@@ -77,6 +77,7 @@
     bank_writeback,
     real_bank_numels,
     register_bank_drivers,
+    set_real_bank_bwd_installed,
 )
 from cppmega_mlx.runtime.path_c_dps_adapter import (
     alloc_device_banks,
@@ -246,6 +247,149 @@ def register_real_forward_driver(leaf, real_driver, bank_numels: dict[str, int],
         tvm_ffi.register_global_func(f"pathc.bank_fwd_{i}", fwd, override=True)
 
 
+# --------------------------------------------------------------------------- #
+# (b2) The REAL-kernel BACKWARD driver — symmetric to make_real_bank_forward_driver.
+#
+# This is what flips the §15-profiled `pathc.bank_bwd_i` from the ABSTRACT NUMPY host
+# backward (path_c_relax_step_banks._region_bwd_driver, the 79.6%/91.4% wall — a host
+# round-trip of the ~2028 MB region banks) to the REAL GRIDDED backward MR kernel
+# (run_backward=1), which behind the call_dps_packed boundary runs the gridded
+# B0/B1/B2 chunked SSD backward on the physical banks. The bank-SSA, remat, optimizer,
+# and loss structure are UNCHANGED — only the per-region backward compute becomes the
+# real device-resident kernel.
+#
+# The bank-backward call_dps_packed ABI (from _region_bwd_driver) is
+#   (actg_in, param, state_ckpt, paramg_in, actg_out, paramg_out):
+#     IN : actg_in   = incoming activation grad (from the next layer)
+#          param     = the model weights (read-only)
+#          state_ckpt= checkpoint i (the forward-saved activation snapshot)
+#          paramg_in = running parameter-grad accumulator
+#     OUT: actg_out  = grad propagated to the previous layer (the bwd kernel writes
+#                      BANK_ACTG in place)
+#          paramg_out= updated parameter-grad accumulator (BANK_PARAMG in place)
+#
+# `leaf` is a BACKWARD PathCRegionLeaf (run_backward=1) carrying the real kernel;
+# `real_driver` is make_real_kernel_driver(leaf, device) from the driver phase.
+# --------------------------------------------------------------------------- #
+def make_real_bank_backward_driver(leaf, real_driver, bank_numels: dict[str, int],
+                                   *, device=None):
+    """Build a `pathc.bank_bwd_*` packed func that runs the REAL gridded MR backward.
+
+    Mirrors make_real_bank_forward_driver exactly (same DEVICE-RESIDENT vs
+    NUMPY-STAGED gate, RULE #1: clear device-vs-host gate, no silent fallback), but
+    over the BACKWARD bank ABI. The real kernel is the SAME compiled artifact as the
+    forward, specialized to run_backward=1 (the gridded B0/B1/B2 transpose). It seeds
+    BANK_ACTG/BANK_PARAM/BANK_STATE/BANK_PARAMG from actg_in/param/state_ckpt/
+    paramg_in, runs the kernel (writes BANK_ACTG + BANK_PARAMG in place), then writes
+    actg_out = BANK_ACTG and paramg_out = BANK_PARAMG.
+
+    RULE #1: no numpy host backward is reachable from this driver; the only paths are
+    (a) DEVICE-RESIDENT real kernel and (b) NUMPY-STAGED real kernel (the same real
+    compiled kernel, just host-staged for the CPU self-check). Any failure RAISES.
+    """
+
+    actg_n = bank_numels[BANK_ACTG]
+    state_n = bank_numels[BANK_STATE]
+    param_n = bank_numels[BANK_PARAM]
+    paramg_n = bank_numels[BANK_PARAMG]
+
+    # Build the device-resident kernel driver + driver-owned device physical banks
+    # ONCE (lazily, first device call); persists across every backward call.
+    dev_state: dict[str, object] = {}
+
+    def _ensure_device_state(dev):
+        if "driver" not in dev_state:
+            dev_state["driver"] = make_device_resident_kernel_driver(leaf, dev)
+            dev_state["banks"] = alloc_device_banks(leaf.bank_shapes, dev)
+        return dev_state["driver"], dev_state["banks"]
+
+    def packed(actg_in, param, state_ckpt, paramg_in, actg_out, paramg_out):
+        if bank_arg_is_device(actg_in):
+            # ---- DEVICE-RESIDENT PATH (no numpy host bounce) ----
+            dev = actg_in.device
+            dev_driver, dbanks = _ensure_device_state(dev)
+            # Seed the backward kernel's physical banks from the SSA banks
+            # (device->device view copies; tail-zeroed where the SSA bank is shorter).
+            bank_copy_prefix_device(
+                dbanks[BANK_ACTG], actg_in,
+                min(actg_n, int(np.prod([int(d) for d in actg_in.shape]))))
+            bank_copy_prefix_device(
+                dbanks[BANK_PARAM], param,
+                min(param_n, int(np.prod([int(d) for d in param.shape]))))
+            bank_copy_prefix_device(
+                dbanks[BANK_STATE], state_ckpt,
+                min(state_n, int(np.prod([int(d) for d in state_ckpt.shape]))))
+            bank_copy_prefix_device(
+                dbanks[BANK_PARAMG], paramg_in,
+                min(paramg_n, int(np.prod([int(d) for d in paramg_in.shape]))))
+            # Run the REAL gridded backward kernel (run_backward=1) on the device
+            # banks (writes BANK_ACTG + BANK_PARAMG in place).
+            dev_driver(leaf, dbanks, dev)
+            # actg_out <- new activation grad (device->device).
+            bank_copy_prefix_device(
+                actg_out, dbanks[BANK_ACTG],
+                min(actg_n, int(np.prod([int(d) for d in actg_out.shape]))))
+            # paramg_out <- updated parameter-grad accumulator (device->device).
+            bank_copy_prefix_device(
+                paramg_out, dbanks[BANK_PARAMG],
+                min(paramg_n, int(np.prod([int(d) for d in paramg_out.shape]))))
+            return
+
+        # ---- NUMPY-STAGED REFERENCE PATH (CPU self-test; STILL the real kernel) ----
+        ag = bank_arg_to_host(actg_in)
+        p = bank_arg_to_host(param)
+        ck = bank_arg_to_host(state_ckpt)
+        pgi = bank_arg_to_host(paramg_in)
+        banks = {
+            BANK_ACT: np.zeros(bank_numels[BANK_ACT], np.float32),
+            BANK_ACTG: np.ascontiguousarray(ag, np.float32).reshape(-1)[:actg_n].copy(),
+            BANK_PARAM: np.ascontiguousarray(p, np.float32).reshape(-1)[:param_n].copy(),
+            BANK_PARAMG: np.ascontiguousarray(pgi, np.float32).reshape(-1)[:paramg_n].copy(),
+            BANK_STATE: np.ascontiguousarray(ck, np.float32).reshape(-1)[:state_n].copy(),
+        }
+        for bank, n in ((BANK_ACTG, actg_n), (BANK_PARAM, param_n),
+                        (BANK_PARAMG, paramg_n), (BANK_STATE, state_n)):
+            if banks[bank].size < n:
+                pad = np.zeros(n, np.float32)
+                pad[: banks[bank].size] = banks[bank]
+                banks[bank] = pad
+        real_driver(leaf, banks)
+        bank_writeback(actg_out, banks[BANK_ACTG].reshape(-1)[:actg_n])
+        bank_writeback(paramg_out, banks[BANK_PARAMG].reshape(-1)[:paramg_n])
+
+    return packed
+
+
+def register_real_backward_driver(leaf, real_driver, bank_numels: dict[str, int],
+                                  n_layers: int, *, device=None) -> None:
+    """Install the REAL gridded path_c backward kernel as EVERY `pathc.bank_bwd_i`.
+
+    This REPLACES the abstract numpy host backward (the §15 79.6%/91.4% wall) with the
+    real device-resident gridded backward MR kernel (run_backward=1 -> gridded
+    B0/B1/B2). It also flips ``set_real_bank_bwd_installed(True)`` so that:
+      * the RULE #1 guard in register_bank_drivers does NOT raise (the gridded path is
+        live, not the forbidden numpy fallback), and
+      * a subsequent register_bank_drivers re-entry does NOT clobber this real bwd
+        binding with the numpy one.
+
+    ``leaf`` MUST be a BACKWARD leaf (run_backward=1); ``device`` enables the
+    DEVICE-RESIDENT backward path (banks stay device tensors, no host bounce). Pass the
+    CUDA device on gb10. RULE #1: a forward leaf here is a wiring bug — RAISE."""
+    if int(getattr(leaf, "run_backward", 0)) != 1:
+        raise RuntimeError(
+            "register_real_backward_driver: leaf.run_backward must be 1 (a BACKWARD "
+            f"leaf), got run_backward={getattr(leaf, 'run_backward', None)!r}. "
+            "Pass the bwd-specialized PathCRegionLeaf (RULE #1: no silent forward "
+            "leaf binding as the backward region)."
+        )
+    bwd = make_real_bank_backward_driver(leaf, real_driver, bank_numels, device=device)
+    for i in range(n_layers):
+        tvm_ffi.register_global_func(f"pathc.bank_bwd_{i}", bwd, override=True)
+    # Mark the real gridded backward as installed so the RULE #1 guard knows the
+    # numpy bank backward is NOT the live path (and is never re-bound over this).
+    set_real_bank_bwd_installed(True)
+
+
 # --------------------------------------------------------------------------- #
 # (c) The whole-step IRModule: fwd (remat) + bwd + in-place Adam + LOSS.
 #
diff --git a/scratch/probe_chunked_backward_cuda_gb10.py b/scratch/probe_chunked_backward_cuda_gb10.py
