feat(relax-cuda): real path_c region runs through call_dps_packed on gb10 CUDA VM

make_real_kernel_driver was PROVEN on Metal; this wires + proves it on tvm.cuda(0)
(gb10, Grace-Blackwell aarch64, CUDA 13.x). Two e2e proofs:

  Stage A (scratch/pr3_cuda_real_kernel_driver_gb10.py): the REAL 17-param MR
  path_c tilelang kernel compiles for target=cuda (127.6s, simplify-gate avoids the
  18-23min mamba3-bwd blowup; 149KB CUDA-C, 1 __shared__) and runs through the
  adapter pack/unpack on tvm.cuda(0): 14,680,063/14,680,064 nonzero (matches the
  Metal 14.68M proof).

  Stage B (scratch/pr3_cuda_relax_vm_gb10.py): the SAME region as an
  R.call_dps_packed leaf in one @R.function, compiled target=cuda, run on
  relax.VirtualMachine(ex, tvm.cuda(0)) with a CUDA device-tensor input. Output on
  cuda:0, 14.68M nonzero, max|VM - direct| = 0.0 (byte-exact). The external-function
  boundary executes the REAL tilelang-CUDA kernel e2e through the Relax graph.

CUDA-specific fix to make_region_dps_packed (path_c_dps_adapter.py): the call_dps_packed
ABI tensors arrive as CUDA device tensors on the CUDA VM; np.from_dlpack RAISES
'Unsupported device in DLTensor' for them. Added device-agnostic
_dps_arg_to_host_numpy (np.from_dlpack for CPU/Metal, .numpy() host-copy for CUDA)
and _dps_writeback_host_to_arg (aliased view for CPU/Metal, host->device copyto for
CUDA). CPU-VM path (path_c_relax_step_real) regression-tested unchanged. Fail-loud:
both paths RAISE if neither readback/writeback route is available. RULE #1 honored.

Author: apstenku123

cppmega_mlx/runtime/path_c_dps_adapter.py

@@ -192,8 +192,13 @@ def _packed(*args: Any) -> None:
                 f"{len(leaf.logical_inputs)+1} args "
                 f"(inputs {len(leaf.logical_inputs)} + 1 output), got {len(args)}"
             )
-        in_arrays = [np.from_dlpack(a) for a in args[:-1]]
-        out_array = np.from_dlpack(args[-1])
+        # Import the call_dps_packed ABI tensors to host numpy. On a CUDA Relax VM
+        # (tvm.cuda(0)) these arrive as device tensors; ``np.from_dlpack`` RAISES
+        # "Unsupported device in DLTensor" for them, so we route device tensors
+        # through ``.numpy()`` (an explicit host copy). CPU/Metal tensors still use
+        # the zero-copy ``np.from_dlpack``. RULE #1: if neither works we RAISE.
+        in_arrays = [_dps_arg_to_host_numpy(a) for a in args[:-1]]
+        out_tensor = args[-1]
 
         # Allocate the physical banks we own.
         banks = {b: np.zeros((n,), dtype=np.float32) for b, n in bank_shapes.items()}
@@ -215,13 +220,67 @@ def _packed(*args: Any) -> None:
         _drive_region_compute(leaf, banks)
 
         # Unpack the logical-output sub-range into the trailing DPS output tensor.
+        # On CUDA the DPS output is a device tensor that does NOT alias a host numpy
+        # view, so we write the result back EXPLICITLY (host->device copy) rather
+        # than via an aliasing ``out_array[...] =``.
         m = lmap[leaf.logical_output]
+        out_host_shape = tuple(int(d) for d in out_tensor.shape)
         out_flat = banks[m.bank][m.offset : m.offset + m.size]
-        out_array[...] = out_flat.reshape(out_array.shape)
+        _dps_writeback_host_to_arg(out_tensor, out_flat.reshape(out_host_shape))
 
     return _packed
 
 
+def _dps_arg_to_host_numpy(arg: Any) -> np.ndarray:
+    """Import a call_dps_packed ABI tensor to a host numpy array, device-agnostically.
+
+    CPU/Metal DLPack-capable tensors import zero-copy via ``np.from_dlpack``; CUDA
+    tensors (whose device DLPack export is rejected by numpy with "Unsupported
+    device in DLTensor") are copied to host via the tensor's ``.numpy()`` method.
+    RULE #1 (fail loud): if BOTH paths fail the original errors are surfaced."""
+
+    try:
+        return np.from_dlpack(arg)
+    except Exception as dlpack_err:  # noqa: BLE001 -- CUDA: Unsupported device in DLTensor
+        to_numpy = getattr(arg, "numpy", None)
+        if to_numpy is None:
+            raise RuntimeError(
+                "FAIL-LOUD: DPS arg is neither host-DLPack-importable nor has a "
+                f".numpy() host-copy method (type={type(arg).__name__}); "
+                f"np.from_dlpack raised: {dlpack_err}"
+            ) from dlpack_err
+        return np.ascontiguousarray(to_numpy())
+
+
+def _dps_writeback_host_to_arg(out_tensor: Any, host_result: np.ndarray) -> None:
+    """Write a host numpy result back into a call_dps_packed output tensor,
+    device-agnostically.
+
+    For a CPU/Metal tensor that aliases a host numpy view we assign through that
+    view; for a CUDA device tensor (which does NOT alias host memory) we build a
+    same-device source tensor and ``copyto`` the device output in place. RULE #1:
+    if neither path is available we RAISE."""
+
+    host_result = np.ascontiguousarray(host_result, dtype=np.float32)
+    # CPU/Metal: alias the output's host buffer and assign in place.
+    try:
+        view = np.from_dlpack(out_tensor)
+        view[...] = host_result.reshape(view.shape)
+        return
+    except Exception:  # noqa: BLE001 -- CUDA: device DLPack rejected by numpy
+        pass
+    # CUDA (and any device tensor with a .device + copyto): host->device copy.
+    dev = getattr(out_tensor, "device", None)
+    copyto = getattr(out_tensor, "copyto", None)
+    if dev is None or copyto is None:
+        raise RuntimeError(
+            "FAIL-LOUD: DPS output tensor is not host-DLPack-aliasable and lacks "
+            f"a (.device, .copyto) device-writeback path (type={type(out_tensor).__name__})"
+        )
+    src = tvm.runtime.tensor(host_result, device=dev)
+    src.copyto(out_tensor)
+
+
 # Hook the region compute. Default = a transparent reference matching the region's
 # logical semantics (so the adapter ABI is testable on CPU without a live GPU); the
 # real path is set by ``set_region_kernel_driver`` to call ``leaf.kernel`` on device.

scratch/pr3_cuda_real_kernel_driver_gb10.py

@@ -0,0 +1,85 @@
+"""Stage A: REAL path_c MR tilelang JITKernel through the call_dps_packed adapter
+pack/unpack on tvm.cuda(0) (gb10). Direct mirror of scratch/pr3_real_kernel_driver.py
+(proven on Metal: 14.68M nonzero) but on CUDA. Proves the on-device real-kernel path
+of the adapter works on gb10 CUDA."""
+from __future__ import annotations
+import sys, time
+import numpy as np
+import tvm, tvm_ffi
+
+from cppmega_mlx.runtime.path_c_dps_adapter import (
+    parse_logical_to_physical, parse_physical_bank_shapes, prim_bank_param_order,
+    PathCRegionLeaf, set_region_kernel_driver, make_region_dps_packed,
+    make_real_kernel_driver,
+)
+from cppmega_mlx.runtime.path_c_fusion import (
+    build_path_c_aot_autograd_region, build_path_c_model_region_from_route_symbols,
+    _compile_tilelang_prim_func, _tilelang_compile_pass_configs_for_prim_func,
+    _path_c_default_target,
+)
+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
+
+
+def build_real_kernel():
+    cfg = local_gb10_quarter_profile().hybrid_config()
+    region = build_path_c_model_region_from_route_symbols(
+        region_name="mr_path_c", route_symbols=("M", "R"), model_config=cfg)
+    prim = path_c_fusion_schedule_template(build_path_c_aot_autograd_region(region))
+    target = _path_c_default_target()
+    assert target == "cuda", f"expected cuda target on gb10, got {target}"
+    t0 = time.time()
+    kernel = _compile_tilelang_prim_func(
+        prim, target=target, execution_backend="tvm_ffi",
+        pass_configs=_tilelang_compile_pass_configs_for_prim_func(prim) or None)
+    src = kernel.get_kernel_source()
+    print(f"  tilelang.compile(target=cuda) OK in {time.time()-t0:.2f}s -> JITKernel, "
+          f"{len(kernel.params)} params, {len(src)} bytes CUDA-C")
+    # report smem usage hint
+    nshared = src.count("__shared__")
+    print(f"  kernel src has {nshared} __shared__ decls")
+    return prim, kernel
+
+
+def main() -> int:
+    print("Stage A -- REAL path_c MR JITKernel through call_dps_packed adapter on tvm.cuda(0)")
+    dev = tvm.cuda(0)
+    print("TVM:", tvm.__version__, "cuda.exist=", dev.exist)
+    if not dev.exist:
+        raise RuntimeError("FAIL: tvm.cuda(0) not present")
+    prim, kernel = build_real_kernel()
+    lmap = parse_logical_to_physical(prim)
+    bank_shapes = parse_physical_bank_shapes(prim)
+    order = prim_bank_param_order(prim)
+    print(f"  ABI: {len(prim.params)} params, {len(lmap)} logical tensors, "
+          f"{len(bank_shapes)} banks; bank order[:5]={order[:5]}")
+
+    leaf = PathCRegionLeaf(
+        name="mr_path_c_fwd", run_backward=0, prim=prim, kernel=kernel,
+        logical_map=lmap, bank_shapes=bank_shapes, bank_param_order=order,
+        logical_inputs=("route_0_M_hidden",), logical_output="route_0_M_hidden_after")
+
+    driver = make_real_kernel_driver(leaf, dev)   # <-- on tvm.cuda(0)
+    set_region_kernel_driver(driver)
+    packed = make_region_dps_packed(leaf)
+
+    m_in = lmap["route_0_M_hidden"]; m_out = lmap["route_0_M_hidden_after"]
+    x = (np.random.rand(*m_in.logical_shape).astype("float32") - 0.5) * 0.01
+    out = np.zeros(m_out.logical_shape, np.float32)
+    print(f"  driving region fwd: in {m_in.logical_shape} -> out {m_out.logical_shape} "
+          f"via REAL CUDA kernel...")
+    t0 = time.time()
+    packed(tvm_ffi.from_dlpack(x), tvm_ffi.from_dlpack(out))
+    dt = time.time()-t0
+    nz = int(np.count_nonzero(out))
+    print(f"  REAL kernel executed on CUDA in {dt:.2f}s; out mean={out.mean():.3e} "
+          f"std={out.std():.3e} nonzero={nz}/{out.size}")
+    if nz == 0:
+        raise RuntimeError("FAIL-LOUD: real CUDA kernel produced ALL-ZERO output")
+    print(f"STAGE-A PASS: REAL path_c MR JITKernel runs on gb10 CUDA THROUGH the "
+          f"call_dps_packed adapter pack/unpack ({nz} nonzero outputs).")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

scratch/pr3_cuda_relax_vm_gb10.py

@@ -0,0 +1,123 @@
+"""Stage B: a REAL path_c MR region as an R.call_dps_packed LEAF inside ONE Relax
+@R.function, COMPILED FOR target=cuda and RUN ON relax.VirtualMachine(ex, tvm.cuda(0)).
+The packed func runs the REAL tilelang-compiled CUDA JITKernel (via make_real_kernel_driver)
+through the adapter pack/unpack. Proves the call_dps_packed external-function boundary
+executes the real path_c-CUDA kernel end-to-end ON gb10 THROUGH the Relax graph + CUDA VM.
+
+Reference check: the SAME real kernel run via the direct adapter driver (Stage A path,
+CPU-imported input) -- the VM graph output must match it (the graph boundary must execute
+the identical real compute, not a stand-in)."""
+from __future__ import annotations
+import sys, time
+import numpy as np
+import tvm, tvm_ffi
+from tvm import relax
+
+from cppmega_mlx.runtime.path_c_dps_adapter import (
+    parse_logical_to_physical, parse_physical_bank_shapes, prim_bank_param_order,
+    PathCRegionLeaf, set_region_kernel_driver, make_region_dps_packed,
+    register_region_dps_packed, make_real_kernel_driver, emit_region_call,
+)
+from cppmega_mlx.runtime.path_c_fusion import (
+    build_path_c_aot_autograd_region, build_path_c_model_region_from_route_symbols,
+    _compile_tilelang_prim_func, _tilelang_compile_pass_configs_for_prim_func,
+    _path_c_default_target,
+)
+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
+
+
+def build_real_kernel():
+    cfg = local_gb10_quarter_profile().hybrid_config()
+    region = build_path_c_model_region_from_route_symbols(
+        region_name="mr_path_c", route_symbols=("M", "R"), model_config=cfg)
+    prim = path_c_fusion_schedule_template(build_path_c_aot_autograd_region(region))
+    target = _path_c_default_target()
+    assert target == "cuda", f"expected cuda target on gb10, got {target}"
+    t0 = time.time()
+    kernel = _compile_tilelang_prim_func(
+        prim, target=target, execution_backend="tvm_ffi",
+        pass_configs=_tilelang_compile_pass_configs_for_prim_func(prim) or None)
+    print(f"  tilelang.compile(target=cuda) OK in {time.time()-t0:.2f}s -> "
+          f"{len(kernel.params)} params, {len(kernel.get_kernel_source())} bytes CUDA-C")
+    return prim, kernel
+
+
+def main() -> int:
+    print("Stage B -- REAL path_c MR call_dps_packed leaf on the CUDA Relax VM (gb10)")
+    dev = tvm.cuda(0)
+    print("TVM:", tvm.__version__, "cuda.exist=", dev.exist)
+    if not dev.exist:
+        raise RuntimeError("FAIL: tvm.cuda(0) not present")
+    prim, kernel = build_real_kernel()
+    lmap = parse_logical_to_physical(prim)
+    bank_shapes = parse_physical_bank_shapes(prim)
+    order = prim_bank_param_order(prim)
+
+    leaf = PathCRegionLeaf(
+        name="mr_path_c_fwd", run_backward=0, prim=prim, kernel=kernel,
+        logical_map=lmap, bank_shapes=bank_shapes, bank_param_order=order,
+        logical_inputs=("route_0_M_hidden",), logical_output="route_0_M_hidden_after")
+
+    # Install the REAL on-device CUDA driver + register the region as a packed func.
+    driver = make_real_kernel_driver(leaf, dev)
+    set_region_kernel_driver(driver)
+    packed_name = "pathc.mr_fwd.cuda"
+    register_region_dps_packed(leaf, packed_name)
+
+    m_in = lmap["route_0_M_hidden"]; m_out = lmap["route_0_M_hidden_after"]
+    in_sinfo = relax.TensorStructInfo(tuple(int(d) for d in m_in.logical_shape), "float32")
+    out_sinfo = relax.TensorStructInfo(tuple(int(d) for d in m_out.logical_shape), "float32")
+
+    # Assemble the @R.function: ONE call_dps_packed leaf (the real region).
+    bb = relax.BlockBuilder()
+    xv = relax.Var("x", in_sinfo)
+    with bb.function("train_step", [xv]):
+        with bb.dataflow():
+            y = emit_region_call(bb, packed_name, [xv], out_sinfo)
+            out = bb.emit_output(y)
+        bb.emit_func_output(out)
+    mod = bb.get()
+    if not relax.analysis.well_formed(mod):
+        raise RuntimeError("FAIL-LOUD: assembled CUDA call_dps_packed step not well-formed")
+    print("  @R.function well_formed: True; compiling for target=cuda + CUDA VM...")
+
+    ex = tvm.compile(mod, target=tvm.target.Target("cuda"))
+    vm = relax.VirtualMachine(ex, dev)
+
+    # Input: a CUDA device tensor fed to the VM (the real graph entry on gb10).
+    x_np = ((np.random.default_rng(0).random(m_in.logical_shape, np.float32) - 0.5)
+            * 0.01).astype(np.float32)
+    x_cuda = tvm.runtime.tensor(x_np, device=dev)
+    print(f"  running CUDA VM train_step: in {m_in.logical_shape} (device={x_cuda.device})...")
+    t0 = time.time()
+    res = vm["train_step"](x_cuda)
+    dev.sync()
+    out_vm = res.numpy()   # CUDA -> host
+    dt = time.time()-t0
+    nz = int(np.count_nonzero(out_vm))
+    print(f"  CUDA VM executed in {dt:.2f}s; out device was {res.device}; "
+          f"mean={out_vm.mean():.3e} std={out_vm.std():.3e} nonzero={nz}/{out_vm.size}")
+    if nz == 0:
+        raise RuntimeError("FAIL-LOUD: CUDA VM call_dps_packed produced ALL-ZERO output")
+
+    # Reference: SAME real kernel via the direct adapter packed (host-imported input).
+    set_region_kernel_driver(make_real_kernel_driver(leaf, dev))
+    direct = make_region_dps_packed(leaf)
+    ref_out = np.zeros(m_out.logical_shape, np.float32)
+    direct(tvm_ffi.from_dlpack(x_np), tvm_ffi.from_dlpack(ref_out))
+    maxdiff = float(np.abs(out_vm - ref_out).max())
+    print(f"  REF (same real CUDA kernel via direct adapter): nonzero="
+          f"{int(np.count_nonzero(ref_out))}; max|VM - REF| = {maxdiff:.3e}")
+    if not np.allclose(out_vm, ref_out, rtol=1e-4, atol=1e-6):
+        raise RuntimeError(f"FAIL-LOUD: CUDA VM graph output disagrees with the same "
+                           f"real kernel run directly; max abs diff={maxdiff}")
+    print(f"STAGE-B PASS: a REAL path_c MR region runs THROUGH the Relax R.call_dps_packed "
+          f"graph on the gb10 CUDA VM ({nz} nonzero, byte-matching the direct real-kernel "
+          f"run, max diff {maxdiff:.1e}). The external-function boundary executes the REAL "
+          f"tilelang-CUDA kernel end-to-end on gb10.")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())