lever dlpack-fix: relieve MLX+torch cudaMallocAsync dual-pool contention at the fp8-backward DLPack bridge

Root-causes the R1-e2e fp8-backward OOM (docs §21): the error
'cudaMallocAsync(&data,size,stream) failed: out of memory' is MLX's OWN
allocator hitting genuine exhaustion of the single GB10 117GB unified pool
because MLX's and torch/TE's stream-ordered mempools each reserve-and-retain
unified memory neither releases (neither sets cudaMemPoolAttrReleaseThreshold).

Python-only, no MLX rebuild required for the primary fix; zero-copy stays the
ONE path (RULE #1 - the bridge still RAISES on a genuinely-uncrossable tensor,
never host-copies or degrades to bf16):
- _cuda_zerocopy.py: ctypes libcudart shim trim_cuda_mempool() sets
  ReleaseThreshold=0 + cudaMemPoolTrimTo(0) on the default device pool;
  relieve_bridge_memory_pressure() also mx.clear_cache()s MLX. Both RAISE on
  CUDA failure. Gated CPPMEGA_FP8_BRIDGE_TRIM (default ON).
- fp8_te_linear.py: _fp8_linear_bwd eval+sync the 3 bridge operands then trims
  at the boundary (after sync, never mid-kernel) before the crossing.
- bench harness: PYTORCH_CUDA_ALLOC_CONF=backend:cudaMallocAsync,release_threshold:0
  set before torch import; fp8 arm sets mx.set_cache_limit(0) +
  torch.cuda.set_per_process_memory_fraction to split the budget.

Author: apstenku123

cppmega_mlx/nn/_tilelang/_cuda_zerocopy.py

@@ -33,6 +33,127 @@
 
 import mlx.core as mx
 
+# ---------------------------------------------------------------------------
+# CUDA mempool release-threshold + trim shim (lever dlpack-fix).
+#
+# Root cause of the R1-e2e fp8-backward OOM (docs/RELAX-GRAPH-VS-MEGATRON.md §21):
+# the error string ``cudaMallocAsync(&data, size, stream) failed: out of memory``
+# is MLX's OWN allocator (mlx/backend/cuda/allocator.cpp:225) hitting genuine
+# physical-memory exhaustion of the single GB10 117 GB unified pool. TWO
+# stream-ordered cudaMallocAsync mempools (MLX's + torch/TE's) each reserve-and-
+# retain unified memory that neither releases, because neither sets
+# ``cudaMemPoolAttrReleaseThreshold`` (MLX only READS ReservedMemCurrent at
+# allocator.cpp:247). The COMBINED reserved set exceeds physical memory.
+#
+# This shim (Python-only, NO MLX rebuild) sets the release threshold to 0 on the
+# device mempool(s) and calls ``cudaMemPoolTrimTo(pool, 0)`` so reserved-but-idle
+# unified memory is RETURNED to the OS before the competing allocation. It is
+# called ONLY at a bridge boundary AFTER eval/sync (never mid-kernel) — see
+# RISK 3 in the lever research. RULE #1: trimming only makes the zero-copy path
+# FIT; it never introduces a host-copy / fp8->bf16 degrade. If libcudart cannot
+# be loaded or a CUDA call fails, this RAISES with where+what (no silent skip) —
+# the caller decides whether trimming is required for its budget.
+# ---------------------------------------------------------------------------
+
+_CUDA_MEMPOOL_ATTR_RELEASE_THRESHOLD = 4  # cudaMemPoolAttrReleaseThreshold
+_cudart = None  # lazily-loaded libcudart handle
+_RELEASE_THRESHOLD_SET: set[int] = set()  # devices whose threshold we already set
+
+
+def _load_cudart():
+    """Load libcudart once; RAISE (RULE #1) with where+what if it cannot load."""
+
+    global _cudart
+    if _cudart is not None:
+        return _cudart
+    last_exc = None
+    for name in ("libcudart.so", "libcudart.so.12", "libcudart.so.13", "libcudart.so.11.0"):
+        try:
+            _cudart = ctypes.CDLL(name)
+            break
+        except OSError as exc:  # noqa: PERF203 - try the next soname
+            last_exc = exc
+    if _cudart is None:
+        raise RuntimeError(
+            "_cuda_zerocopy: cannot load libcudart for the CUDA mempool trim shim "
+            f"(tried libcudart.so / .so.12 / .so.13 / .so.11.0): {last_exc}. The "
+            "fp8 bridge needs cudaMemPoolTrimTo to fit the dual-pool budget."
+        )
+    # cudaGetDevice(int*), cudaDeviceGetDefaultMemPool(pool*, int dev),
+    # cudaMemPoolSetAttribute(pool, attr, void*), cudaMemPoolTrimTo(pool, size_t)
+    _cudart.cudaGetDevice.argtypes = [ctypes.POINTER(ctypes.c_int)]
+    _cudart.cudaGetDevice.restype = ctypes.c_int
+    _cudart.cudaDeviceGetDefaultMemPool.argtypes = [
+        ctypes.POINTER(ctypes.c_void_p),
+        ctypes.c_int,
+    ]
+    _cudart.cudaDeviceGetDefaultMemPool.restype = ctypes.c_int
+    _cudart.cudaMemPoolSetAttribute.argtypes = [
+        ctypes.c_void_p,
+        ctypes.c_int,
+        ctypes.c_void_p,
+    ]
+    _cudart.cudaMemPoolSetAttribute.restype = ctypes.c_int
+    _cudart.cudaMemPoolTrimTo.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
+    _cudart.cudaMemPoolTrimTo.restype = ctypes.c_int
+    return _cudart
+
+
+def _default_mempool(rt, device: int) -> ctypes.c_void_p:
+    pool = ctypes.c_void_p()
+    err = rt.cudaDeviceGetDefaultMemPool(ctypes.byref(pool), int(device))
+    if err != 0:
+        raise RuntimeError(
+            f"_cuda_zerocopy: cudaDeviceGetDefaultMemPool(dev={device}) failed "
+            f"(cudaError={err}); cannot trim the CUDA mempool for the fp8 bridge."
+        )
+    return pool
+
+
+def trim_cuda_mempool(device: int | None = None) -> None:
+    """Set ReleaseThreshold=0 and trim the default mempool to 0 on ``device``.
+
+    Called at the fp8 bridge boundary (after eval/sync) to return reserved-but-idle
+    unified memory so the competing MLX+torch+TE allocations fit the single GB10
+    pool. RAISES (RULE #1) on any CUDA failure — never silently degrades. If the
+    box genuinely has no idle reserved memory to trim this is a fast no-op (TrimTo
+    returns cudaSuccess and reclaims nothing).
+    """
+
+    rt = _load_cudart()
+    if device is None:
+        dev = ctypes.c_int(0)
+        err = rt.cudaGetDevice(ctypes.byref(dev))
+        if err != 0:
+            raise RuntimeError(
+                f"_cuda_zerocopy: cudaGetDevice failed (cudaError={err}) resolving "
+                "the device for the CUDA mempool trim."
+            )
+        device = int(dev.value)
+    pool = _default_mempool(rt, device)
+    # Set the release threshold to 0 ONCE per device so the driver does not retain
+    # reserved memory across the trim (default threshold is UINT64_MAX = retain all).
+    if device not in _RELEASE_THRESHOLD_SET:
+        thresh = ctypes.c_uint64(0)
+        err = rt.cudaMemPoolSetAttribute(
+            pool,
+            _CUDA_MEMPOOL_ATTR_RELEASE_THRESHOLD,
+            ctypes.cast(ctypes.byref(thresh), ctypes.c_void_p),
+        )
+        if err != 0:
+            raise RuntimeError(
+                f"_cuda_zerocopy: cudaMemPoolSetAttribute(ReleaseThreshold=0, "
+                f"dev={device}) failed (cudaError={err}); cannot bound the CUDA "
+                "mempool reservation for the fp8 bridge."
+            )
+        _RELEASE_THRESHOLD_SET.add(device)
+    err = rt.cudaMemPoolTrimTo(pool, ctypes.c_size_t(0))
+    if err != 0:
+        raise RuntimeError(
+            f"_cuda_zerocopy: cudaMemPoolTrimTo(dev={device}, 0) failed "
+            f"(cudaError={err}); cannot return idle reserved unified memory."
+        )
+
 # ---------------------------------------------------------------------------
 # DLPack C ABI (DLPack >= 0.5; matches tvm-ffi's expected layout)
 # ---------------------------------------------------------------------------
@@ -116,6 +237,32 @@ def zerocopy_enabled() -> bool:
     return os.environ.get("CPPMEGA_TILELANG_CUDA_ZEROCOPY", "0") not in ("", "0", "false", "False")
 
 
+def relieve_bridge_memory_pressure(device: int | None = None) -> None:
+    """Trim BOTH the MLX cache and the CUDA mempool at a bridge boundary.
+
+    The fp8 backward crosses the MLX<->torch DLPack bridge while MLX's forward/scan
+    working set and torch/TE's fp8 scratch both hold reserved unified memory. This
+    helper, called immediately BEFORE a backward crossing (after the producer
+    eval/sync), returns MLX's cached buffers (``mx.clear_cache``) and returns the
+    CUDA mempool's reserved-but-idle unified memory (``trim_cuda_mempool``) so the
+    next allocation fits the single GB10 pool. It is gated by
+    ``CPPMEGA_FP8_BRIDGE_TRIM`` (default ON when unset, to fix the §21 OOM; set 0 to
+    A/B the un-trimmed behavior). RULE #1: this only makes the zero-copy path FIT —
+    it never host-copies or degrades. On a genuine CUDA failure it RAISES with
+    where+what (no silent skip).
+    """
+
+    if os.environ.get("CPPMEGA_FP8_BRIDGE_TRIM", "1").strip().lower() in ("0", "false", "no", "off"):
+        return
+    # MLX side: return cached (non-live) buffers so the unified pool shrinks.
+    clear = getattr(mx, "clear_cache", None)
+    if callable(clear):
+        clear()
+    # CUDA side: set ReleaseThreshold=0 + cudaMemPoolTrimTo(0) so the driver
+    # returns reserved-but-idle unified memory shared with torch/TE's pool.
+    trim_cuda_mempool(device)
+
+
 def _dlpack_device_type() -> int:
     """Resolve the DLPack device_type to emit (kDLCUDA(2) default; 13 for A/B)."""
 

cppmega_mlx/nn/_tilelang/fp8_te_linear.py

@@ -294,6 +294,7 @@ def _fp8_linear_bwd(
 
     from cppmega_mlx.nn._tilelang._cuda_zerocopy import (
         mlx_cuda_array_to_torch_tensor,
+        relieve_bridge_memory_pressure,
         torch_cuda_tensor_to_mlx,
     )
 
@@ -313,6 +314,26 @@ def _fp8_linear_bwd(
     x2 = x.reshape((-1, k_in)) if x.ndim != 2 else x
     g2 = cotangent.reshape((-1, n_out)) if cotangent.ndim != 2 else cotangent
 
+    # Materialize the three bridge operands and flush MLX's stream BEFORE trimming.
+    # trim_cuda_mempool must run only at a bridge boundary after eval/sync (RISK 3:
+    # trimming mid-kernel could reclaim memory live tensors depend on). After this
+    # eval, x2/w/g2 are the live set; the cache + idle reserved memory are not.
+    mx.eval(x2, weight, g2)
+    mx.synchronize()
+    # Relieve the dual-pool (MLX + torch/TE) cudaMallocAsync reservation contention
+    # that causes the §21 fp8-backward OOM, so the zero-copy crossing FITS. RULE #1:
+    # this does not host-copy — it returns idle reserved unified memory and the
+    # bridge below still RAISES on a genuinely-uncrossable tensor.
+    try:
+        relieve_bridge_memory_pressure()
+    except Exception as exc:  # noqa: BLE001 - RULE #1 surface where+what
+        raise RuntimeError(
+            f"fp8_te_linear(bwd): CUDA mempool trim before the fp8-backward bridge "
+            f"crossing failed (site {site_key!r}): {type(exc).__name__}: {exc}. The "
+            "zero-copy bridge cannot be made to fit; this is a hard memory boundary, "
+            "not a fallback point."
+        ) from exc
+
     try:
         x_t = mlx_cuda_array_to_torch_tensor(x2)
         w_t = mlx_cuda_array_to_torch_tensor(weight)

scripts/bench_fp8_e2e_step_20260604.py

@@ -36,6 +36,21 @@
 import statistics
 import time
 
+# RULE #1 (lever dlpack-fix): bound torch/TE's cudaMallocAsync pool so it does NOT
+# reserve-and-retain unified memory that contends with MLX's pool on the single
+# GB10 117 GB unified part (the §21 fp8-backward OOM root cause). This env MUST be
+# set BEFORE the first ``import torch`` anywhere in the process; torch reads it once
+# at CUDA-context init. release_threshold:0 makes torch's stream-ordered pool return
+# reserved-but-idle memory instead of hoarding it, so the combined MLX+torch+TE
+# reserved set fits. Overridable via CPPMEGA_TORCH_ALLOC_CONF for A/B. This is NOT a
+# fallback — it only makes the zero-copy fp8 bridge FIT; the bridge still RAISES on a
+# genuinely-uncrossable tensor.
+if "PYTORCH_CUDA_ALLOC_CONF" not in os.environ:
+    os.environ["PYTORCH_CUDA_ALLOC_CONF"] = os.environ.get(
+        "CPPMEGA_TORCH_ALLOC_CONF",
+        "backend:cudaMallocAsync,release_threshold:0",
+    )
+
 # RULE #1: fix the NVRTC loader BEFORE torch / TE are imported anywhere (gb10).
 # Safe no-op off-gb10. Must precede the first transformer_engine import (which
 # happens lazily inside fp8_te_linear when the fp8 arm runs).
@@ -87,6 +102,40 @@ def _torch_reset_peak() -> None:
         pass
 
 
+def _budget_mlx_torch() -> None:
+    """Split the single GB10 unified pool between MLX and torch/TE for the fp8 arm.
+
+    lever dlpack-fix: the §21 fp8-backward OOM is dual-pool reservation contention,
+    not activation size. Three knobs, all bounds (NOT fallbacks):
+      (1) ``mx.set_cache_limit(0)`` — MLX returns freed buffers immediately so the
+          bridge trim (``relieve_bridge_memory_pressure``) can reclaim them, instead
+          of MLX hoarding them in its cache.
+      (2) ``torch.cuda.set_per_process_memory_fraction(F)`` — caps torch/TE's share
+          so MLX + torch/TE + the transient per-tensor export doubling sum < the
+          physical 117 GB part. F is set conservatively (env CPPMEGA_TORCH_MEM_FRAC,
+          default 0.45 = ~52 GB on a 117 GB box) leaving MLX the majority.
+      (3) ``mx.set_memory_limit`` honored via CPPMEGA_MLX_MEMORY_LIMIT_GB (in main).
+    RULE #1: if the split is too tight the bridge still RAISES with where+what — it
+    never host-copies or degrades to bf16.
+    """
+
+    set_cache = getattr(mx, "set_cache_limit", None)
+    if callable(set_cache):
+        set_cache(0)
+    try:
+        import torch
+
+        if torch.cuda.is_available():
+            frac = float(os.environ.get("CPPMEGA_TORCH_MEM_FRAC", "0.45"))
+            torch.cuda.set_per_process_memory_fraction(frac)
+    except Exception as exc:  # RULE #1: a budget-cap failure is a config bug, surface it.
+        raise RuntimeError(
+            "bench_fp8_e2e: could not apply the torch CUDA memory-fraction budget for "
+            f"the fp8 arm (CPPMEGA_TORCH_MEM_FRAC): {type(exc).__name__}: {exc}. The "
+            "fp8 bridge needs the MLX+torch budget split to fit the GB10 unified pool."
+        ) from exc
+
+
 def _run_arm(arm: str, args, tokens) -> dict:
     """Run one arm (bf16 | fp8) for warmup+steps timed steps; return the metrics."""
 
@@ -99,6 +148,12 @@ def _run_arm(arm: str, args, tokens) -> dict:
             os.environ["CPPMEGA_FP8_LINEAR_MIN_K"] = str(args.fp8_floor)
             os.environ["CPPMEGA_FP8_LINEAR_MIN_N"] = str(args.fp8_floor)
             os.environ["CPPMEGA_FP8_LINEAR_MIN_M"] = str(args.fp8_floor)
+        # lever dlpack-fix: split the single GB10 unified pool between MLX and
+        # torch/TE so neither hoards. (1) Stop MLX from caching freed buffers so the
+        # bridge trim can return them. (2) Cap torch's share so MLX + torch/TE + the
+        # transient export doubling sum < physical. These are bounds, not fallbacks:
+        # if the split is too tight the bridge RAISES (no host-copy degrade).
+        _budget_mlx_torch()
     else:
         os.environ.pop(FP8_LINEAR_ENV, None)