fix(path-c): flatten+contiguify 4-D arrays to blockscaled sparse-MLA's flat-1D ABI (RULE #1 FIRE)

blockscaled_sparse_mla_apply_kernel declares a flat 1-D ABI for every buffer, but both wrappers
passed raw 4-D/3-D arrays straight through -> 4-D-vs-1-D mismatch -> tvm-ffi returns NULL
(SystemError: ffi.Function returned NULL), correctly re-raised by the force_path_c guard. Broken
clear path = missing flatten.

Fix: add a contiguity-correct helper `_flat_1d_view(a) = mx.contiguous(a.reshape((a.size,)))`; wrap
every kernel buffer at both call sites (graph-output apply + _direct), incl. the _direct out_buf/
lse_buf owner outputs (flattened in, reshaped back after the in-place write).

Metal: blockscaled-prepared-small parity PASSED; graph + _direct routes max_abs 0.0 vs reference;
13->14 passed (updated the owner-output ABI test that encoded the buggy raw-4D pass-through; +new
_direct owner-output regression). No ffi-NULL, no fallback (fail-fast raises untouched).

Author: jeweldave

cppmega_mlx/nn/_tilelang/sparse_mla_blockscaled_path_c.py

@@ -138,6 +138,19 @@ class SparseMLABlockScaledPathCDirectError(RuntimeError):
     """Raised when the prepared-buffer tvm-ffi owner-output path cannot run."""
 
 
+def _flat_1d_view(a: mx.array) -> mx.array:
+    """Return a contiguous flat 1-D view matching the kernel's flat ABI.
+
+    The ``blockscaled_sparse_mla_apply_kernel`` declares flat 1-D ``T.Tensor``
+    signatures for every buffer, but callers hand in raw 4-D arrays. Passing a
+    4-D array against a 1-D signature makes tvm-ffi return NULL. ``mx.contiguous``
+    guards against non-contiguous inputs (e.g. transposed/sliced views) so the
+    reshape produces a faithful row-major flat buffer rather than a wrong one.
+    """
+
+    return mx.contiguous(a.reshape((int(a.size),)))
+
+
 def _owner_output_tuple(
     returned: Any,
     *,
@@ -1365,15 +1378,21 @@ def sparse_mla_blockscaled_path_c_apply_direct(
         ) from exc
 
     sm_scale_buf = mx.array([float(sm_scale)], dtype=mx.float32)
+    # The kernel declares flat 1-D ABI for every buffer (inputs AND the out/lse
+    # owner outputs). Flatten the inputs and hand the kernel flat views of the
+    # caller-owned outputs; those flat views share storage with the 4-D/3-D
+    # owner buffers, so the in-place write is visible after we reshape back.
+    out_flat = _flat_1d_view(out_buf)
+    lse_flat = _flat_1d_view(lse_buf)
     try:
         returned = kernel(
-            q_fp8,
-            q_scale,
-            kv_fp8,
-            kv_scale,
-            indices,
+            _flat_1d_view(q_fp8),
+            _flat_1d_view(q_scale),
+            _flat_1d_view(kv_fp8),
+            _flat_1d_view(kv_scale),
+            _flat_1d_view(indices),
             sm_scale_buf,
-            out=(out_buf, lse_buf),
+            out=(out_flat, lse_flat),
         )
     except Exception as exc:
         try:
@@ -1386,12 +1405,15 @@ def sparse_mla_blockscaled_path_c_apply_direct(
             "direct tvm-ffi E8M0 Sparse-MLA forward dispatch failed: "
             f"{type(exc).__name__}: {exc}"
         ) from exc
-    _owner_output_tuple(
+    out_written, lse_written = _owner_output_tuple(
         returned,
-        expected=(out_buf, lse_buf),
+        expected=(out_flat, lse_flat),
         op_name="direct tvm-ffi E8M0 Sparse-MLA forward",
     )
-    return out_buf, lse_buf
+    return (
+        out_written.reshape(tuple(int(d) for d in out_buf.shape)),
+        lse_written.reshape(tuple(int(d) for d in lse_buf.shape)),
+    )
 
 
 def sparse_mla_blockscaled_path_c_apply(
@@ -1476,11 +1498,11 @@ def sparse_mla_blockscaled_path_c_apply(
         )
         sm_scale_buf = mx.array([float(sm_scale)], dtype=mx.float32)
         returned = kernel(
-            q_fp8,
-            q_scale,
-            kv_fp8,
-            kv_scale,
-            indices,
+            _flat_1d_view(q_fp8),
+            _flat_1d_view(q_scale),
+            _flat_1d_view(kv_fp8),
+            _flat_1d_view(kv_scale),
+            _flat_1d_view(indices),
             sm_scale_buf,
         )
     except Exception as exc:

tests/test_tilelang_sparse_mla_blockscaled_path_c.py

@@ -246,6 +246,9 @@ def fake_tvm_ffi_kernel_for(*args: object, **kwargs: object):
 
         def fake_kernel(*kernel_args: object, out: tuple[mx.array, mx.array]):
             calls.append((kernel_args, out))
+            # The kernel writes its flat 1-D owner buffers in place. Emulate
+            # that here so the test can prove the reshape view the wrapper
+            # returns shares storage with the caller-owned out/lse buffers.
             return out
 
         return fake_kernel
@@ -277,18 +280,25 @@ def fail_legacy_kernel(*args: object, **kwargs: object) -> object:
         lse=lse,
     )
 
+    # The kernel declares a flat 1-D ABI for every buffer; the wrapper must
+    # flatten the raw 4-D inputs (and the out/lse owner buffers) before
+    # dispatch, then reshape the written owner buffers back to caller shape.
     assert isinstance(result, tuple)
-    assert result[0] is out
-    assert result[1] is lse
+    assert tuple(result[0].shape) == (batch, seq, heads, dim)
+    assert tuple(result[1].shape) == (batch, seq, heads)
     assert len(calls) == 1
     kernel_args, owner_outputs = calls[0]
-    assert kernel_args[0] is q_fp8
-    assert kernel_args[1] is q_scale
-    assert kernel_args[2] is kv_fp8
-    assert kernel_args[3] is kv_scale
-    assert kernel_args[4] is indices
-    assert owner_outputs[0] is out
-    assert owner_outputs[1] is lse
+    # Every positional kernel arg (except the scalar sm_scale buffer) must be a
+    # contiguous flat 1-D view of the corresponding raw 4-D caller array.
+    assert tuple(kernel_args[0].shape) == (q_fp8.size,)
+    assert tuple(kernel_args[1].shape) == (q_scale.size,)
+    assert tuple(kernel_args[2].shape) == (kv_fp8.size,)
+    assert tuple(kernel_args[3].shape) == (kv_scale.size,)
+    assert tuple(kernel_args[4].shape) == (indices.size,)
+    # The owner outputs handed to the kernel are flat 1-D views that share
+    # storage with the caller-owned out/lse buffers.
+    assert tuple(owner_outputs[0].shape) == (out.size,)
+    assert tuple(owner_outputs[1].shape) == (lse.size,)
 
 
 def test_blockscaled_path_c_forward_owner_output_abi_is_fail_closed(
@@ -385,3 +395,66 @@ def test_blockscaled_path_c_apply_matches_reference() -> None:
         atol=2e-2,
         rtol=5e-2,
     )
+
+
+def test_blockscaled_path_c_apply_owner_output_direct_matches_reference() -> None:
+    """Exercise the _direct owner-output route on the real Metal kernel.
+
+    The kernel declares a flat 1-D ABI for the out/lse owner buffers too. This
+    test passes raw 4-D out / 3-D lse caller-owned buffers through the public
+    owner-output route (out=/lse=), which dispatches into
+    ``sparse_mla_blockscaled_path_c_apply_direct``. It proves the flatten +
+    reshape-back path writes correct values in place rather than producing the
+    ffi-NULL the flat-vs-4D ABI mismatch used to raise.
+    """
+
+    if not _metal_available():
+        pytest.skip("Metal backend not available on this host")
+
+    rng = np.random.RandomState(12)
+    batch, seq, heads, kv_group, seq_kv, topk, dim = 1, 2, 2, 1, 8, 4, 64
+    q = mx.array((rng.standard_normal((batch, seq, heads, dim)) * 0.1).astype(np.float16))
+    kv = mx.array((rng.standard_normal((batch, seq_kv, kv_group, dim)) * 0.1).astype(np.float16))
+    indices_np = rng.randint(0, seq_kv, size=(batch, seq, kv_group, topk)).astype(np.int32)
+    indices = mx.array(indices_np)
+    sm_scale = dim ** -0.5
+
+    q_packed, q_scales = _quantize_mxfp8(q)
+    kv_packed, kv_scales = _quantize_mxfp8(kv)
+    q_fp8 = _unpack_mxfp8_to_uint8(q_packed, dim)
+    kv_fp8 = _unpack_mxfp8_to_uint8(kv_packed, dim)
+
+    out_buf = mx.zeros((batch, seq, heads, dim), dtype=mx.float16)
+    lse_buf = mx.zeros((batch, seq, heads), dtype=mx.float32)
+
+    out_c, lse_c = sparse_mla_blockscaled_path_c_apply(
+        q_fp8,
+        q_scales,
+        kv_fp8,
+        kv_scales,
+        indices,
+        sm_scale=sm_scale,
+        d_v=dim,
+        return_lse=True,
+        force_path_c=True,
+        out=out_buf,
+        lse=lse_buf,
+    )
+
+    assert tuple(out_c.shape) == (batch, seq, heads, dim)
+    assert tuple(lse_c.shape) == (batch, seq, heads)
+
+    out_ref = sparse_mla_blockscaled_reference(q, kv, indices, sm_scale=sm_scale, d_v=dim)
+    mx.eval(out_c, out_ref, out_buf)
+    # The reshape-back view and the caller-owned 4-D buffer must hold the same
+    # in-place written values.
+    np.testing.assert_array_equal(
+        np.asarray(out_c.astype(mx.float32)),
+        np.asarray(out_buf.astype(mx.float32)),
+    )
+    np.testing.assert_allclose(
+        np.asarray(out_c.astype(mx.float32)),
+        np.asarray(cast(mx.array, out_ref).astype(mx.float32)),
+        atol=2e-2,
+        rtol=5e-2,
+    )