fix(path-c): materialize contiguous tvm-ffi inputs for mamba3 (DLPack non-contiguous)

RULE #1 surfaced this: mamba3 Path C passed caller MLX arrays straight across the tvm-ffi boundary
where dlpack_to_tvm_tensor rejects non-contiguous layouts (FromDLPack: Tensor is not contiguous).
Root cause: x,B,C,z arrive non-contiguous from the producer (reshape of a strided _split_by_sizes
slice; mean + RoPE-state-dim-transpose + group->head-broadcast); A/dt/D/h0 contiguous.

Fix at the Path C ABI source (no fallback): _materialize_contiguous_inputs runs mx.contiguous on
each tvm-ffi input (layout-only, dtype-preserving -- unsupported dtypes still raise; no-op on
already-contiguous) at all 3 dispatch entries (fwd, fwd_with_snapshots, bwd Metal entry).

Metal: all 5 DLPack-contiguity failures cleared; +regression test (passes here, fails on pristine).
NOTE: the remaining 44 mamba3 Path C failures are the SAME ffi-NULL (ffi.Function returned NULL)
common blocker as GDN/KDA Path C + #13 -- the tvm-ffi launch-recovery covers only the 5-param fp8
signature, not the mamba3/GDN/KDA/#13 kernel signatures. That is the next (highest-leverage) fix.

Author: jeweldave

cppmega_mlx/nn/_tilelang/mamba3_path_c.py

@@ -2356,6 +2356,34 @@ def _require_supported_no_hidden_casts(
     return dtypes
 
 
+def _materialize_contiguous_inputs(
+    op_name: str,
+    *named_arrays: tuple[str, mx.array],
+) -> tuple[mx.array, ...]:
+    """Materialize each tvm-ffi input as a contiguous, DLPack-exportable buffer.
+
+    Path C dispatches caller-owned MLX arrays straight across the tvm-ffi
+    boundary, where ``dlpack_to_tvm_tensor`` rejects any non-contiguous layout
+    with ``FromDLPack: Tensor is not contiguous``. The Mamba3 producer (the
+    model layer feeding ``_dispatch_mamba3_scan``) can legitimately hand Path C
+    strided/broadcast views — e.g. RoPE state-dim transposes/reshapes or
+    group->head broadcasts upstream of the scan. Those are *layout* views of an
+    otherwise supported buffer, so the correct fix is to materialize the
+    contiguous buffer the ABI requires at this choke point, not to fall back to
+    another path (RULE #1).
+
+    This is layout-only: ``mx.contiguous`` is a no-op on already-contiguous
+    arrays and never changes dtype, so the no-hidden-cast contract is preserved
+    (unsupported dtypes are still rejected by
+    :func:`_require_supported_no_hidden_casts`). Anything ``mx.contiguous``
+    genuinely cannot make DLPack-exportable still surfaces loudly at the
+    dispatch boundary via :func:`_raise_if_dlpack_boundary_failure`.
+    """
+
+    del op_name
+    return tuple(mx.contiguous(array) for _name, array in named_arrays)
+
+
 def _require_owner_array(
     op_name: str,
     name: str,
@@ -2770,6 +2798,21 @@ def mamba3_mimo_fwd_path_c(
         ("D", D),
         ("h0", h0),
     )
+    # The producer may hand Path C strided/broadcast views (RoPE state-dim
+    # transposes, group->head broadcasts, slices). The tvm-ffi DLPack import
+    # rejects non-contiguous buffers, so materialize the contiguous layout the
+    # ABI requires here (layout-only, dtype-preserving; RULE #1: no fallback).
+    x, B, C, z, A, dt, D, h0 = _materialize_contiguous_inputs(
+        "mamba3_mimo_fwd_path_c",
+        ("x", x),
+        ("B", B),
+        ("C", C),
+        ("z", z),
+        ("A", A),
+        ("dt", dt),
+        ("D", D),
+        ("h0", h0),
+    )
     if seq == 0:
         if out is not None:
             raise RuntimeError(
@@ -2878,6 +2921,20 @@ def _mamba3_mimo_fwd_path_c_with_snapshots(
         ("D", D),
         ("h0", h0),
     )
+    # See mamba3_mimo_fwd_path_c: materialize contiguous tvm-ffi inputs so a
+    # strided/broadcast producer view does not trip the DLPack contiguity check
+    # (layout-only, dtype-preserving; RULE #1: no fallback).
+    x, B, C, z, A, dt, D, h0 = _materialize_contiguous_inputs(
+        "mamba3_mimo_fwd_path_c",
+        ("x", x),
+        ("B", B),
+        ("C", C),
+        ("z", z),
+        ("A", A),
+        ("dt", dt),
+        ("D", D),
+        ("h0", h0),
+    )
     if seq == 0:
         raise RuntimeError(
             "mamba3_mimo_fwd_path_c snapshot route is not dispatchable for "
@@ -3860,6 +3917,23 @@ def _ref(x_, B_, C_, z_, A_, dt_, D_, h0_):
         _y, vjps = mx.vjp(_ref, primals, [dy])
         return tuple(vjps)
 
+    # Materialize contiguous tvm-ffi inputs once at the Metal backward entry so
+    # every downstream Path C bwd kernel (snapshot, lane-grad, reducer) receives
+    # DLPack-exportable buffers even when the VJP cotangent or a producer view
+    # is strided/broadcast (layout-only, dtype-preserving; RULE #1: no
+    # fallback). Unsupported dtypes are still rejected downstream.
+    dy, x, B, C, z, A, dt, D, h0 = _materialize_contiguous_inputs(
+        "mamba3_mimo_bwd_path_c",
+        ("dy", dy),
+        ("x", x),
+        ("B", B),
+        ("C", C),
+        ("z", z),
+        ("A", A),
+        ("dt", dt),
+        ("D", D),
+        ("h0", h0),
+    )
     return _mamba3_mimo_bwd_path_c_kernel(dy, x, B, C, z, A, dt, D, h0, out=out)
 
 

tests/test_tilelang_mamba3_path_c.py

@@ -202,6 +202,59 @@ def _require_mamba3_path_c() -> None:
         pytest.skip(f"mamba3 Path C unavailable on this host: {status.reason}")
 
 
+def _is_dlpack_contiguity_failure(exc: BaseException) -> bool:
+    """True iff any link of the cause/context chain is a DLPack contiguity reject."""
+
+    seen: set[int] = set()
+    cur: BaseException | None = exc
+    while cur is not None and id(cur) not in seen:
+        seen.add(id(cur))
+        text = str(cur)
+        if "not contiguous" in text or "DLPackConversionError" in type(cur).__name__:
+            return True
+        cur = cur.__cause__ or cur.__context__
+    return False
+
+
+def test_fwd_path_c_accepts_strided_producer_views_no_dlpack_error() -> None:
+    """Strided/broadcast producer views must not trip the tvm-ffi DLPack check.
+
+    The Mamba3 model layer hands Path C non-contiguous views (RoPE state-dim
+    transposes, group->head broadcasts, reshapes of sliced projections). Path C
+    materializes them contiguous at its ABI boundary, so dispatch must never
+    raise ``FromDLPack: Tensor is not contiguous`` (RULE #1: fix at the source,
+    no fallback). Any *other* dispatch failure (e.g. a tilelang build-level
+    ffi-NULL) is out of scope for this contiguity guard and tolerated here.
+    """
+
+    _require_mamba3_path_c()
+    batch, seq, heads, headdim, state = 1, 4, 2, 4, 5
+    mx.random.seed(0)
+
+    def rnd(*shape: int) -> mx.array:
+        return mx.random.normal(shape).astype(mx.float32)
+
+    # h0 as a transpose view (strided) and B as a broadcast view (stride-0).
+    x = rnd(batch, seq, heads, headdim)
+    B = mx.broadcast_to(rnd(batch, seq, 1, state), (batch, seq, heads, state))
+    C = rnd(batch, seq, heads, state)
+    z = rnd(batch, seq, heads, headdim)
+    A = rnd(batch, seq, heads)
+    dt = mx.abs(rnd(batch, seq, heads))
+    D = rnd(heads)
+    h0 = mx.transpose(rnd(batch, heads, state, headdim), (0, 1, 3, 2))
+    mx.eval(x, B, C, z, A, dt, D, h0)
+
+    try:
+        y, h_last = mamba3_path_c.mamba3_mimo_fwd_path_c(x, B, C, z, A, dt, D, h0)
+        mx.eval(y, h_last)
+    except Exception as exc:  # noqa: BLE001 - we classify and re-assert below
+        assert not _is_dlpack_contiguity_failure(exc), (
+            "Path C still rejects strided/broadcast inputs at the DLPack "
+            f"boundary: {exc}"
+        )
+
+
 def _assert_no_bwd_lane_grad_outputs(msl: str) -> None:
     for name in _BWD_LANE_GRAD_OUTPUT_TOKENS:
         assert name not in msl