feat(tilelang-cuda): zero-copy torch->MLX writeback for fused path_c (no host bounce)

The path_c fused sparse-MLA EAGER kernels feed inputs to torch zero-copy via
DLPack (MLX-CUDA kDLCUDA export), but the kernel OUTPUT (torch CUDA tensor) was
written back into MLX with a host bounce: `_cuda_eager._torch_cuda_to_mlx` did
`t.detach().cpu().numpy()` because MLX could not IMPORT a CUDA DLPack capsule.
That .cpu() PCIe roundtrip erased the fused kernel's latency win (kernel-bound ->
overhead-bound), even though the result was numerically correct and memory-
efficient.

MLX now imports CUDA DLPack natively (DatasunriseOU/mlx upstream-integration:
cuda_dlpack_to_mlx + cu::import_external_buffer), so this closes the output-side
gap symmetrically:

- _cuda_zerocopy.py: new `torch_cuda_tensor_to_mlx(t, out_dtype)` imports a torch
  CUDA tensor into an mx.array via DLPack with NO host roundtrip (torch
  __dlpack__ consumed by the native MLX CUDA import; a single device-side copy
  lands it in an MLX-owned buffer). Plus `mlx_cuda_import_available()` capability
  probe. RULE #1: raises (never silently host-bounces) if the tensor is not CUDA
  or the MLX build cannot import the capsule.
- _cuda_eager.py: `_torch_cuda_to_mlx` now uses the zero-copy writeback when
  `CPPMEGA_TILELANG_CUDA_ZEROCOPY=1` (same gate as the input bridge); raises if
  the build can't import. With the env unset it uses the explicit eager numpy-
  host copy (a deliberately-selected mode, not a silent degrade). Module docstring
  updated.
- _sparse_mla_v32_fused.py: docstring updated — the fused fwd/bwd now stays
  GPU-resident on both the input and output boundary, keeping the kernel latency
  win.

Verified on gb10 / sm_121 (GB10): torch CUDA -> mx.array round-trips bit-
identical (f32/f16/bf16/int32), the fused fwd runs correct through the zero-copy
writeback.

Author: apstenku123

cppmega_mlx/nn/_tilelang/_cuda_eager.py

@@ -23,13 +23,16 @@
   VJP (documented TODO to port the Metal SIMD lane-grad bwd).
 * **m2rnn** — not yet ported; see :func:`m2rnn_supported_cuda_eager` (TODO).
 
-MLX on this host has no CUDA-array DLPack export
-(``a.__dlpack__`` raises "CUDA DLPack export is not supported"), so the
-MLX<->kernel boundary uses a numpy host roundtrip into torch CUDA tensors.
-That is correct (if not zero-copy) for an EAGER fallback whose job is to run a
-real TileLang-CUDA kernel rather than raise. The Metal path is untouched: every
-public entry below is only invoked from a ``cuda_eager_available()`` branch the
-callers add *after* the existing ``can_run_metal()`` checks.
+The MLX<->kernel boundary is zero-copy on this fork. Inputs are exported to torch
+zero-copy via DLPack (MLX-CUDA kDLCUDA export, commit 6da6a0e4), and kernel
+outputs are imported back into MLX zero-copy via DLPack (MLX-CUDA kDLCUDA
+*import*: ``cuda_dlpack_to_mlx``) when ``CPPMEGA_TILELANG_CUDA_ZEROCOPY=1`` —
+no ``.cpu()`` host roundtrip, so the fused kernel keeps its latency win. With the
+env unset the writeback falls back to the explicit eager numpy-host copy (a
+deliberately-selected mode, not a silent degrade); see ``_torch_cuda_to_mlx``.
+The Metal path is untouched: every public entry below is only invoked from a
+``cuda_eager_available()`` branch the callers add *after* the existing
+``can_run_metal()`` checks.
 """
 
 # pyright: reportMissingImports=false, reportInvalidTypeForm=false
@@ -141,12 +144,40 @@ def _mlx_to_torch_cuda(a: mx.array) -> Any:
     return t.contiguous()
 
 
+def _zerocopy_writeback_enabled() -> bool:
+    """Whether the kernel-output writeback should use the zero-copy CUDA import.
+
+    Same env gate as the input bridge (``CPPMEGA_TILELANG_CUDA_ZEROCOPY``): when
+    set, the torch-CUDA kernel result is imported back into MLX via DLPack with
+    no ``.cpu()`` host roundtrip (DatasunriseOU native MLX CUDA import). When
+    unset (default) the explicit eager numpy-host copy below is used. These are
+    two deliberately-selected modes, NOT a silent fallback.
+    """
+
+    from cppmega_mlx.nn._tilelang._cuda_zerocopy import zerocopy_enabled
+
+    return zerocopy_enabled()
+
+
 def _torch_cuda_to_mlx(t: Any, out_dtype: Any) -> mx.array:
-    """Copy a CUDA torch tensor back into an ``mx.array`` of ``out_dtype``."""
+    """Move a CUDA torch tensor back into an ``mx.array`` of ``out_dtype``.
+
+    With ``CPPMEGA_TILELANG_CUDA_ZEROCOPY=1`` this stays GPU-resident: the result
+    crosses torch->MLX via DLPack (no ``.cpu()`` / PCIe bounce), which is what
+    preserves the fused kernel's latency win. RULE #1: when the zero-copy mode is
+    requested it RAISES if the MLX build cannot import the CUDA capsule, instead
+    of silently host-bouncing. Default (env unset) uses the eager numpy-host copy.
+    """
 
-    import numpy as np  # noqa: F401
     import torch
 
+    if _zerocopy_writeback_enabled():
+        from cppmega_mlx.nn._tilelang._cuda_zerocopy import torch_cuda_tensor_to_mlx
+
+        return torch_cuda_tensor_to_mlx(t, out_dtype)
+
+    import numpy as np  # noqa: F401
+
     cpu = t.detach().to(device="cpu")
     if cpu.dtype == torch.bfloat16:
         cpu = cpu.to(torch.float32)

cppmega_mlx/nn/_tilelang/_cuda_zerocopy.py

@@ -374,3 +374,99 @@ def mlx_cuda_array_to_torch_tensor(arr: "mx.array") -> Any:
         )
     del torch
     return t
+
+
+# ---------------------------------------------------------------------------
+# OUTPUT side: torch CUDA result -> mx.array, zero-copy, no host bounce
+# ---------------------------------------------------------------------------
+#
+# Symmetric counterpart to ``mlx_cuda_array_to_torch_tensor`` (the INPUT bridge).
+# A TileLang/torch-CUDA kernel writes its result into a torch CUDA tensor; this
+# imports it straight back into an ``mx.array`` via DLPack with NO ``.cpu()``
+# host roundtrip. It relies on the native MLX CUDA *import* (DatasunriseOU fork:
+# convert.cpp ``cuda_dlpack_to_mlx`` + cuda backend ``import_external_buffer``),
+# which wraps the foreign CUDA buffer zero-copy and materializes it into an
+# MLX-owned GPU allocation with a single device-side copy (no PCIe bounce).
+#
+# RULE #1: if the running MLX build lacks the CUDA import (older binary that
+# still raises "CUDA DLPack import is not supported"), we RAISE with where+what
+# instead of silently host-bouncing. The caller (``_cuda_eager``) chooses this
+# zero-copy writeback only when the build supports it.
+
+
+def mlx_cuda_import_available() -> bool:
+    """Return True iff the running MLX build can import a CUDA DLPack capsule.
+
+    Probes the native import by round-tripping a tiny torch CUDA tensor through
+    ``mx.array(t)`` (DLPack). Returns False (not raise) for the *capability*
+    probe; actual writebacks raise on failure.
+    """
+
+    try:
+        import torch
+    except Exception:
+        return False
+    if not torch.cuda.is_available():
+        return False
+    try:
+        t = torch.ones(4, dtype=torch.float32, device="cuda")
+        torch.cuda.synchronize()
+        a = mx.array(t)
+        mx.eval(a)
+        return bool(abs(float(a.sum().item()) - 4.0) < 1e-4)
+    except Exception:
+        return False
+
+
+def torch_cuda_tensor_to_mlx(t, out_dtype=None):
+    """Import a torch CUDA tensor into an ``mx.array`` zero-copy (no host bounce).
+
+    Replaces the ``t.detach().cpu().numpy()`` writeback in
+    ``_cuda_eager._torch_cuda_to_mlx``: the kernel result stays GPU-resident and
+    crosses the torch->MLX boundary via DLPack (torch ``__dlpack__`` consumed by
+    the native MLX CUDA import). A single device-side copy lands it in an
+    MLX-owned buffer; there is no ``.cpu()`` / PCIe roundtrip.
+
+    RAISES (RULE #1) if the tensor is not a CUDA tensor or the MLX build cannot
+    import the CUDA DLPack capsule -- never silently degrades to a host copy.
+    """
+
+    import torch
+
+    if not isinstance(t, torch.Tensor):
+        raise TypeError(
+            f"_cuda_zerocopy.torch_cuda_tensor_to_mlx: expected a torch.Tensor, "
+            f"got {type(t).__name__}."
+        )
+    if not t.is_cuda:
+        raise RuntimeError(
+            f"_cuda_zerocopy.torch_cuda_tensor_to_mlx: tensor is on {t.device}, "
+            f"not CUDA; the zero-copy MLX import requires a CUDA device tensor."
+        )
+
+    # DLPack producer-readiness contract: flush torch's CUDA stream so the
+    # buffer is materialized before MLX reads it. MLX import does a full device
+    # sync of its own before the device-side copy, but we sync the producer here
+    # to honor the handoff ordering explicitly.
+    src = t.detach().contiguous()
+    torch.cuda.synchronize()
+
+    try:
+        # mx.array(torch_tensor) routes through MLX create_array ->
+        # nd_array_to_mlx -> (DatasunriseOU) cuda_dlpack_to_mlx, consuming the
+        # tensor's __dlpack__ kDLCUDA capsule with no host roundtrip.
+        arr = mx.array(src)
+    except Exception as exc:  # noqa: BLE001 - surface where + what failed
+        raise RuntimeError(
+            f"_cuda_zerocopy.torch_cuda_tensor_to_mlx: mx.array() rejected the "
+            f"torch CUDA DLPack capsule (dtype={src.dtype}, "
+            f"shape={tuple(int(d) for d in src.shape)}). This MLX build likely "
+            f"lacks the native CUDA DLPack import (cuda_dlpack_to_mlx); rebuild "
+            f"MLX-CUDA with the import patch. Underlying error: "
+            f"{type(exc).__name__}: {exc}"
+        ) from exc
+
+    if out_dtype is not None and arr.dtype != out_dtype:
+        arr = arr.astype(out_dtype)
+    del torch
+    return arr

cppmega_mlx/nn/_tilelang/_sparse_mla_v32_fused.py

@@ -41,9 +41,14 @@
 sm_12x and the original Hopper kernel on sm_90/sm_100. MLX-CUDA q/kv/indices are
 fed to the torch-backend kernel interfaces zero-copy via DLPack when
 ``CPPMEGA_TILELANG_CUDA_ZEROCOPY=1`` (kDLCUDA ``DLManagedTensor`` over
-``mx.array.data_ptr()``, no host roundtrip); outputs come back from torch (MLX
-cannot import a CUDA DLPack, so the torch->MLX writeback is a host bounce — an
-MLX limitation, surfaced honestly, not a silent degrade of the input path).
+``mx.array.data_ptr()``, no host roundtrip). The kernel outputs are imported
+back from torch into MLX zero-copy too (DatasunriseOU MLX-CUDA DLPack *import*:
+``cuda_dlpack_to_mlx`` wraps the foreign CUDA buffer and materializes it into an
+MLX-owned GPU allocation with a single device-side copy — no ``.cpu()`` host
+bounce), so the whole fused fwd/bwd stays GPU-resident and keeps the kernel's
+latency win. With ``CPPMEGA_TILELANG_CUDA_ZEROCOPY`` unset the writeback uses the
+explicit eager numpy-host copy (a deliberately-selected mode, not a silent
+degrade).
 """
 
 from __future__ import annotations