feat(path_c): device-resident DPS driver — eliminate the per-step host numpy bounce

Rework the physical-bank DPS driver so the banks stay tvm.cuda(0) device
nd-arrays end-to-end in the forward, removing the numpy host round-trip that
was ~96.9% of the Relax-graph train_step wall (docs/RELAX-GRAPH-VS-MEGATRON §4).

path_c_dps_adapter.py:
  - device_bank_view / device_pack / device_unpack: zero-copy device sub-range
    slicing via Tensor._create_view(byte_offset) + device->device copyto (no numpy).
  - alloc_device_banks: banks via tvm.runtime.empty on device, zeroed once, reused.
  - make_device_resident_kernel_driver: builds the kernel arg list once (5 device
    banks + curried run_backward + device-resident zero scratch), per call runs
    kernel(*args)+sync; kernel mutates banks in place at out_idx (nothing to read back).
  - make_region_dps_packed._packed routes by output DLPack device type
    (__dlpack_device__()[0] != kDLCPU): device VM -> device-resident path, CPU VM ->
    numpy reference (clear device-vs-host gate, not a try/except fallback).
  - make_real_kernel_driver kept as the numpy-staged reference (CPU self-test + equiv).

path_c_relax_step_banks.py:
  - bank_arg_is_device (DLPack-device gate at the VM boundary), bank_copy_prefix_device
    (device->device prefix copy + device tail-zero), _zero_device_tensor (device-zero
    buffer cache; only first fill per (shape,dtype,device) touches host).

path_c_relax_train_step.py:
  - make_real_bank_forward_driver(device=dev): device-resident forward — seeds the
    kernel act/param banks via device->device view copies, runs the device-resident
    kernel in place, fills act_out/state_out device-side. No .numpy() at the VM
    boundary in the forward.

Validated on gb10 tvm.cuda(0), REAL MR JITKernel
(scratch/pr7_device_resident_driver_validate_gb10.py):
  (A) numeric equivalence device-resident vs numpy-staged: max|dev-numpy|=1.025e-06.
  (B) whole train_step runs e2e on CUDA, loss finite (5.525e-06, 2L x 3 steps).
  (C) per-forward-call 8472ms -> 6515ms = 1.30x (matches scout ceiling for this
      single-launch MR kernel; the larger lever is the removed VM-boundary host
      round-trip of the full 2028 MB banks per forward+remat region per step).

bwd/adam/loss remain the abstract numpy drivers (lever 4); forward+remat is now
device-resident. Forward/remat region fusion deferred.

Author: apstenku123

cppmega_mlx/runtime/path_c_dps_adapter.py

@@ -141,8 +141,32 @@ def sinfo(self) -> relax.TensorStructInfo:
 # (host copy) on read and tvm.runtime.tensor(...).copyto(out) (host->device) on write.
 # RULE #1 (fail loud): if neither route exists we RAISE; no silent fallback.
 # --------------------------------------------------------------------------- #
+def bank_arg_is_device(arg) -> bool:
+    """True iff a call_dps_packed ABI bank tensor lives on a real (non-CPU) device and
+    supports the zero-copy device view + device->device copy ABI. On a CUDA/Metal Relax
+    VM the bank tensors arrive as ``tvm.runtime.Tensor`` (with ``_create_view`` and
+    ``copyto``) on the device -- so the DEVICE-RESIDENT path applies and NO ``.numpy()``
+    host bounce is needed. This is the gate that removes the doc's 96.9% host round-trip
+    at the VM boundary (a clear device-vs-host gate, NOT a try/except fallback)."""
+    if not (hasattr(arg, "_create_view") and hasattr(arg, "copyto")):
+        return False
+    fn = getattr(arg, "__dlpack_device__", None)  # (device_type, device_id)
+    if fn is None:
+        return False
+    try:
+        dt = int(fn()[0])
+    except Exception:  # noqa: BLE001
+        return False
+    return dt != 1  # DLDeviceType.kDLCPU == 1; CUDA=2, Metal=8 -> device-resident path
+
+
 def bank_arg_to_host(arg) -> np.ndarray:
-    """Import a call_dps_packed ABI tensor to host numpy, device-agnostically."""
+    """Import a call_dps_packed ABI tensor to host numpy, device-agnostically.
+
+    NOTE: this is the HOST path -- it forces a device->host copy on CUDA. The
+    device-resident drivers route bank tensors through the device-view helpers above
+    instead and never hit this. It remains for the abstract reference drivers + CPU VM
+    self-test (and for places where a host numpy view is genuinely required)."""
     try:
         return np.from_dlpack(arg)
     except Exception as dlpack_err:  # noqa: BLE001 -- CUDA: Unsupported device in DLTensor
@@ -155,6 +179,51 @@ def bank_arg_to_host(arg) -> np.ndarray:
         return np.ascontiguousarray(to_numpy())
 
 
+def bank_copy_prefix_device(dst, src, n: int) -> None:
+    """Device->device copy of the first ``n`` elements of ``src`` into ``dst`` (both
+    DEVICE tvm.runtime.Tensor of dtype float32), then ZERO the tail of ``dst`` if it is
+    longer. ZERO host traffic: a device view copy + (when needed) one device zero-fill.
+    Mirrors the numpy ``dst[:n] = src[:n]; dst[n:] = 0`` used by the abstract drivers,
+    but device-resident. RULE #1: out-of-range RAISES."""
+    src_n = int(np.prod([int(d) for d in src.shape]))
+    dst_n = int(np.prod([int(d) for d in dst.shape]))
+    if n > src_n or n > dst_n:
+        raise RuntimeError(
+            f"FAIL-LOUD: bank_copy_prefix_device n={n} exceeds src={src_n}/dst={dst_n}")
+    if dst_n > n:
+        _zero_device_tensor(dst)  # zero the tail; the prefix is overwritten next
+    src_v = src._create_view((n,), "float32", relative_byte_offset=0)
+    dst_v = dst._create_view((n,), "float32", relative_byte_offset=0)
+    src_v.copyto(dst_v)
+
+
+_BANK_ZERO_HOST_CACHE: dict[tuple, np.ndarray] = {}
+_BANK_ZERO_DEVICE_CACHE: dict[tuple, object] = {}
+
+
+def _zero_device_tensor(t) -> None:
+    """Zero a DEVICE tensor in place WITHOUT a per-call host->device transfer: copy from
+    a cached DEVICE-resident zero buffer (one per (shape,dtype,device), created once via
+    a single host->device fill). Subsequent zeroings are device->device copies, so the
+    inner loop has NO host traffic. Falls back to a host-zero copyfrom only the FIRST
+    time a given (shape,dtype,device) is seen (to materialise the device-zero buffer)."""
+    shp = tuple(int(d) for d in t.shape)
+    dtype = str(t.dtype)
+    dev = t.device
+    dkey = (shp, dtype, repr(dev))  # repr stable: device(type='cuda', index=0)
+    zdev = _BANK_ZERO_DEVICE_CACHE.get(dkey)
+    if zdev is None:
+        hkey = (shp, dtype)
+        zh = _BANK_ZERO_HOST_CACHE.get(hkey)
+        if zh is None:
+            zh = np.zeros(shp, dtype=np.dtype(dtype))
+            _BANK_ZERO_HOST_CACHE[hkey] = zh
+        zdev = tvm.runtime.empty(shp, dtype, device=dev)
+        zdev.copyfrom(zh)  # ONE host->device fill per (shape,dtype,device)
+        _BANK_ZERO_DEVICE_CACHE[dkey] = zdev
+    zdev.copyto(t)  # device->device zero (no host traffic)
+
+
 def bank_writeback(out_tensor, host_result: np.ndarray) -> None:
     """Write a host numpy result into a call_dps_packed output tensor, device-agnostic."""
     host_result = np.ascontiguousarray(host_result, np.float32)

cppmega_mlx/runtime/path_c_relax_step_banks.py

@@ -71,11 +71,17 @@
     BANK_PARAMG,
     BANK_STATE,
     BankSinfo,
+    bank_arg_is_device,
     bank_arg_to_host,
+    bank_copy_prefix_device,
     bank_writeback,
     real_bank_numels,
     register_bank_drivers,
 )
+from cppmega_mlx.runtime.path_c_dps_adapter import (
+    alloc_device_banks,
+    make_device_resident_kernel_driver,
+)
 from cppmega_mlx.runtime.path_c_relax_step_remat import (
     checkpoint_boundaries,
     nearest_boundary,
@@ -126,26 +132,78 @@ def _loss_reference(act_bank: np.ndarray) -> float:
 # `leaf` is a fwd-only PathCRegionLeaf (run_backward=0) carrying the real kernel;
 # `real_driver` is make_real_kernel_driver(leaf, device) from the driver phase.
 # --------------------------------------------------------------------------- #
-def make_real_bank_forward_driver(leaf, real_driver, bank_numels: dict[str, int]):
+def make_real_bank_forward_driver(leaf, real_driver, bank_numels: dict[str, int],
+                                  *, device=None):
     """Build a `pathc.bank_fwd_*` packed func that runs the REAL path_c CUDA kernel.
 
-    ABI: (act_in, param, state_in, act_out, state_out). It packs the act/param banks
-    into the real kernel's physical banks, runs the real kernel (real_driver writes
-    the activation bank back into ``banks``), then writes act_out = activation bank
-    and state_out = checkpoint snapshot of the activation. The real kernel is the
-    SAME compiled artifact proven on gb10 (14.68M nonzero)."""
+    ABI: (act_in, param, state_in, act_out, state_out). It seeds the real kernel's
+    activation/parameter physical banks from the incoming act/param SSA bank tensors,
+    runs the real kernel (writes the activation bank in place), then writes
+    act_out = activation bank and state_out = checkpoint snapshot of the activation.
+    The real kernel is the SAME compiled artifact proven on gb10 (14.68M nonzero).
+
+    TWO PATHS (RULE #1: clear device-vs-host gate, no silent fallback):
+      * DEVICE-RESIDENT (gb10): when the call_dps_packed bank tensors arrive on a real
+        device, the act/param banks are seeded into the driver-OWNED device physical
+        banks via DEVICE->DEVICE view copies, the DEVICE-RESIDENT kernel driver runs
+        the kernel on those device banks IN PLACE, and act_out/state_out are filled by
+        DEVICE->DEVICE copies. NO ``.numpy()`` / host staging anywhere -- this removes
+        the per-region per-step ~2028 MB host round-trip at the VM boundary.
+      * NUMPY-STAGED (CPU self-test / reference): host numpy banks + the numpy-staged
+        real_driver. Used only when the ABI tensors are host tensors.
+
+    ``device`` is required for the device-resident path; the device kernel driver +
+    device physical banks are built ONCE here and reused across calls (the kernel
+    mutates them in place)."""
 
     act_n = bank_numels[BANK_ACT]
     state_n = bank_numels[BANK_STATE]
     param_n = bank_numels[BANK_PARAM]
 
+    # Build the device-resident kernel driver + the driver-owned device physical banks
+    # ONCE (lazily, on first device call). These persist across every forward 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(act_in, param, state_in, act_out, state_out):
+        if bank_arg_is_device(act_in):
+            # ---- DEVICE-RESIDENT PATH (no numpy host bounce) ----
+            dev = act_in.device
+            dev_driver, dbanks = _ensure_device_state(dev)
+            # The non-input banks (ACTG, PARAMG, STATE) are zeroed ONCE at alloc (in
+            # alloc_device_banks). The scout validated (pr7_device_resident_kernel_feed:
+            # max|device-numpy|=0) that REUSING the banks across calls -- WITHOUT
+            # re-zeroing -- is bit-identical to the fresh-zeroed-numpy path for the
+            # forward kernel (the forward reads ACT+PARAM and overwrites its out banks;
+            # it does not read stale non-input bank state). So we DO NOT re-zero per
+            # call -- that would reintroduce a ~1.5 GB host->device bounce we are here
+            # to eliminate. (Numeric equivalence is asserted on gb10.)
+            # Seed the kernel's activation/parameter physical banks from the SSA banks
+            # (device->device view copies; tail-zeroed where the SSA bank is shorter).
+            bank_copy_prefix_device(dbanks[BANK_ACT], act_in,
+                                    min(act_n, int(np.prod([int(d) for d in act_in.shape]))))
+            bank_copy_prefix_device(dbanks[BANK_PARAM], param,
+                                    min(param_n, int(np.prod([int(d) for d in param.shape]))))
+            # Run the REAL kernel on the device banks (writes BANK_ACT in place).
+            dev_driver(leaf, dbanks, dev)
+            # act_out <- new activation bank (device->device).
+            bank_copy_prefix_device(act_out, dbanks[BANK_ACT],
+                                    min(act_n, int(np.prod([int(d) for d in act_out.shape]))))
+            # state_out <- checkpoint snapshot of the activation (device->device).
+            so_n = int(np.prod([int(d) for d in state_out.shape]))
+            bank_copy_prefix_device(state_out, dbanks[BANK_ACT], min(so_n, act_n))
+            return
+
+        # ---- NUMPY-STAGED REFERENCE PATH (CPU self-test) ----
         a = bank_arg_to_host(act_in)
         p = bank_arg_to_host(param)
         ao = act_out
         so = state_out
-        # Physical banks the real kernel operates on. The activation bank is seeded
-        # with the incoming activation; the parameter bank with the incoming params.
         banks = {
             BANK_ACT: np.ascontiguousarray(a, np.float32).reshape(-1)[:act_n].copy(),
             BANK_ACTG: np.zeros(bank_numels[BANK_ACTG], np.float32),
@@ -161,12 +219,9 @@ def packed(act_in, param, state_in, act_out, state_out):
             pad = np.zeros(param_n, np.float32)
             pad[: banks[BANK_PARAM].size] = banks[BANK_PARAM]
             banks[BANK_PARAM] = pad
-        # Run the REAL tilelang path_c CUDA kernel on the banks (writes act bank).
         real_driver(leaf, banks)
         act_new = banks[BANK_ACT].reshape(-1)
-        # act_out <- new activation bank.
         bank_writeback(ao, act_new[:act_n])
-        # state_out <- checkpoint snapshot of the activation (what backward reads).
         ck = np.zeros(state_n, np.float32)
         n = min(state_n, act_new.size)
         ck[:n] = act_new[:n]
@@ -176,15 +231,17 @@ def packed(act_in, param, state_in, act_out, state_out):
 
 
 def register_real_forward_driver(leaf, real_driver, bank_numels: dict[str, int],
-                                 n_layers: int) -> None:
+                                 n_layers: int, *, device=None) -> None:
     """Install the REAL path_c CUDA kernel as EVERY `pathc.bank_fwd_i` region.
 
     All forward regions share the one real MR kernel (the bank ABI is per-block
     identical; the bank STORAGE is what the planner threads). The backward + optimizer
     + loss drivers stay the abstract numpy ones (the backward kernel is the s_tir wall
     documented in the adapter; the forward proves the real-kernel-through-graph path).
-    """
-    fwd = make_real_bank_forward_driver(leaf, real_driver, bank_numels)
+
+    ``device`` enables the DEVICE-RESIDENT forward path (banks stay device tensors, no
+    host bounce). Pass the CUDA device on gb10."""
+    fwd = make_real_bank_forward_driver(leaf, real_driver, bank_numels, device=device)
     for i in range(n_layers):
         tvm_ffi.register_global_func(f"pathc.bank_fwd_{i}", fwd, override=True)