fix(path-c): close topk_selector masked-interval scaffold + fix owner-output route bug

topk_selector masked intervals fell back to the pure-MLX reference. Two blockers (NOT a
codegen gap - the masked prim topk_selector_kernel already applies the mask in-kernel via
valid=(j>=s)&(j<e)):
- an artificial guard raised TopKPathCDirectError for any masked call;
- a pre-existing, mask-independent bug: the `returned is not indices` identity check broke
  the ENTIRE owner-output route (masked + unmasked) on the installed TileLang, because
  tvm-ffi out_idx=1 returns a fresh MLX wrapper around the caller buffer, not the literal
  object (confirmed failing on pristine origin/main).

Fix: remove the masked guard; relax the ownership check to validate + mirror the returned
int32 (B,k) wrapper into the caller-owned buffer; pass starts/ends through AUTO + explicit
dispatch.

Metal: masked AUTO parity 10/10 set-equal to reference (5 shapes x fp32/fp16), all via Path C
(0 ref fallbacks); short/empty -> -1 sentinels. No irreducible sub-case. No TileLang C++ change
/ no rebuild. Tests 96 passed (1 pre-existing brittle string-match fail confirmed on pristine);
Path A/B + static-interval path unchanged.

Author: jeweldave

cppmega_mlx/nn/_tilelang/topk_selector.py

@@ -60,12 +60,15 @@
 Performance vs MLX's built-in ``argpartition``
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
-MLX's ``mx.argpartition`` is itself an optimised Metal kernel. Unmasked AUTO
-routes now use the TileLang/tvm-ffi owner-output path when it is available;
-masked intervals fall back to the pure-MLX reference until the owner-output
-Path C route grows masked-row parity. Explicit ``backend="metal"`` is retained
-only as a fail-closed compatibility spelling for callers that still probe the
-retired Path B surface.
+MLX's ``mx.argpartition`` is itself an optimised Metal kernel. Both unmasked
+and masked AUTO routes now use the TileLang/tvm-ffi owner-output path when it is
+available: the Path C prim reads ``starts``/``ends`` and masks columns outside
+``[start, end)`` to ``-inf`` before the per-lane top-K insertion, so masked rows
+reach set-equality parity with the pure-MLX reference (including the ``-1``
+sentinel for intervals shorter than ``k``). The pure-MLX reference is used only
+as a fallback when Path C cannot dispatch. Explicit ``backend="metal"`` is
+retained only as a fail-closed compatibility spelling for callers that still
+probe the retired Path B surface.
 
 Algorithm sketch::
 
@@ -725,12 +728,12 @@ def topk_selector_tilelang_direct(
         raise ValueError(
             f"topk_selector expects a (B, T) array; got shape {scores.shape}"
         )
-    if starts is not None or ends is not None:
-        raise TopKPathCDirectError(
-            "topk_selector_tilelang direct tvm-ffi route is currently "
-            "limited to unmasked rows; masked start/end intervals remain on "
-            "Path B until the TileLang owner-output lowering is parity-clean"
-        )
+    # Masked start/end intervals are handled directly by the Path C prim:
+    # the kernel reads ``starts[bx]``/``ends[bx]`` and masks columns outside
+    # ``[start, end)`` to ``-inf`` before the per-lane top-K insertion, so the
+    # owner-output route is parity-clean for masked rows (set-equality with the
+    # pure-MLX reference, including the ``-1`` sentinel when an interval yields
+    # fewer than ``k`` valid columns). No separate masked prim is required.
     batch = int(scores.shape[0])
     seq_len = int(scores.shape[1])
     if k < 1 or k > seq_len:
@@ -779,14 +782,29 @@ def topk_selector_tilelang_direct(
         raise TopKPathCDirectError(
             f"direct tvm-ffi topk dispatch failed: {type(exc).__name__}: {exc}"
         ) from exc
-    if returned is not indices:
-        raise TopKPathCDirectError(
-            "direct tvm-ffi topk did not return the caller-owned output"
-        )
     # TVM encodes into MLX's current Metal command buffer outside the MLX graph.
     # The returned owner array has no lazy MLX dependency edge, so force the
     # command buffer to complete before Python or downstream consumers observe it.
     mx.synchronize()
+    if returned is indices:
+        # tvm-ffi handed back the exact caller-owned object.
+        return indices
+    # Some installed TileLang/tvm-ffi builds return a fresh MLX wrapper around
+    # the caller's ``out_idx`` buffer instead of the literal Python object. The
+    # kernel still wrote into the caller's storage (out_idx=1 == ``indices``),
+    # so the ownership contract holds. Validate the wrapper is a compatible
+    # int32 (batch, k) array and mirror its values into the caller-owned buffer
+    # so downstream consumers see the result on the exact object they passed.
+    if (
+        not isinstance(returned, mx.array)
+        or tuple(returned.shape) != tuple(indices.shape)
+        or returned.dtype != indices.dtype
+    ):
+        raise TopKPathCDirectError(
+            "direct tvm-ffi topk did not return the caller-owned output"
+        )
+    indices[:, :] = returned
+    mx.eval(indices)
     return indices
 
 
@@ -843,10 +861,9 @@ class PathBStatus:
 
 
 _PATH_B_RETIRED_REASON = (
-    "topk_selector direct-MSL Path B is retired. Unmasked AUTO calls should use "
-    "the TileLang/tvm-ffi owner-output route; masked or unsupported no-output "
-    "calls fall back to the pure-MLX reference until masked owner-output Path C "
-    "parity lands."
+    "topk_selector direct-MSL Path B is retired. Both unmasked and masked AUTO "
+    "calls use the TileLang/tvm-ffi owner-output route; only calls where Path C "
+    "cannot dispatch fall back to the pure-MLX reference."
 )
 
 
@@ -881,11 +898,11 @@ def topk_selector(
         starts: optional (B,) int32 starts for per-row mask.
         ends: optional (B,) int32 ends for per-row mask.
         backend: ``"auto"`` (default) tries TileLang Path C with an owner-output
-            buffer for unmasked calls, then falls back to the pure-MLX reference;
-            ``"metal"`` fails closed because Path B is retired;
-            ``"tilelang"`` / ``"path_c"`` require the owner-output Path C
-            helper and allocate the public result buffer for unmasked calls;
-            ``"mlx"`` always uses the reference.
+            buffer for both unmasked and masked (start/end interval) calls, then
+            falls back to the pure-MLX reference; ``"metal"`` fails closed because
+            Path B is retired; ``"tilelang"`` / ``"path_c"`` require the
+            owner-output Path C helper and allocate the public result buffer for
+            unmasked and masked calls; ``"mlx"`` always uses the reference.
     """
 
     if backend not in {"auto", "mlx", "metal", "tilelang", "path_c"}:
@@ -897,15 +914,21 @@ def topk_selector(
     if backend in {"tilelang", "path_c"}:
         try:
             owner_out = mx.zeros((int(scores.shape[0]), int(k)), dtype=mx.int32)
-            return topk_selector_tilelang_direct(scores, k, out=owner_out)
+            return topk_selector_tilelang_direct(
+                scores, k, starts=starts, ends=ends, out=owner_out
+            )
         except TopKPathCDirectError as exc:
             raise RuntimeError("topk_selector: TileLang Path C path unavailable") from exc
-    if starts is None and ends is None:
-        try:
-            owner_out = mx.zeros((int(scores.shape[0]), int(k)), dtype=mx.int32)
-            return topk_selector_tilelang_direct(scores, k, out=owner_out)
-        except TopKPathCDirectError:
-            pass
+    # AUTO routes both unmasked and masked intervals through the owner-output
+    # Path C prim (the kernel applies the [start, end) mask itself); any
+    # dispatch/compile failure falls back to the pure-MLX reference.
+    try:
+        owner_out = mx.zeros((int(scores.shape[0]), int(k)), dtype=mx.int32)
+        return topk_selector_tilelang_direct(
+            scores, k, starts=starts, ends=ends, out=owner_out
+        )
+    except TopKPathCDirectError:
+        pass
     return topk_selector_reference(scores, k, starts=starts, ends=ends)
 
 

tests/test_tilelang_topk.py

@@ -537,9 +537,10 @@ def test_path_c_direct_tvm_ffi_reuses_owner_output_and_mutates_buffer() -> None:
     assert _to_index_sets(out) == _to_index_sets(out_ref)
 
 
-def test_path_c_direct_rejects_starts_ends_mask(
-    monkeypatch: pytest.MonkeyPatch,
-) -> None:
+def test_path_c_direct_masked_interval_parity_set_equality() -> None:
+    """Masked start/end intervals now go through the owner-output Path C prim
+    (no fallback to the pure-MLX reference) and match it for set membership."""
+
     rng = np.random.default_rng(20260504)
     batch, seq_len, k = 3, 64, 4
     scores_np = rng.standard_normal((batch, seq_len)).astype(np.float32)
@@ -552,45 +553,50 @@ def test_path_c_direct_rejects_starts_ends_mask(
     scores = mx.array(scores_np)
     starts = mx.array(starts_np)
     ends = mx.array(ends_np)
-    out = mx.full((batch, k), -123, dtype=mx.int32)
 
-    monkeypatch.setattr(
-        topk_selector_mod,
-        "_path_c_tvm_ffi_kernel_for",
-        lambda *_, **__: (_ for _ in ()).throw(
-            AssertionError("masked direct Path C must fail before compile")
-        ),
-    )
+    out_c = _topk_tilelang_direct_output(scores, k, starts=starts, ends=ends)
+    out_ref = topk_selector_reference(scores, k, starts=starts, ends=ends)
+    mx.eval(out_c, out_ref)
 
-    with pytest.raises(TopKPathCDirectError, match="limited to unmasked rows"):
-        topk_selector_tilelang_direct(scores, k, starts=starts, ends=ends, out=out)
-    mx.eval(out)
-    assert np.asarray(out).tolist() == [[-123] * k] * batch
+    assert _to_index_sets(out_c) == _to_index_sets(out_ref)
+    assert out_c.dtype == mx.int32
+    assert tuple(out_c.shape) == (batch, k)
+    # Every selected index lies inside its row's [start, end) interval.
+    arr = np.asarray(out_c)
+    for b in range(batch):
+        for idx in arr[b]:
+            if idx >= 0:
+                assert starts_np[b] <= idx < ends_np[b]
 
 
 @pytest.mark.parametrize(
     "starts_np,ends_np",
     [
         (np.array([0, 3], dtype=np.int32), None),
         (None, np.array([-1, 99], dtype=np.int32)),
+        (np.array([1, 2], dtype=np.int32), np.array([3, 2], dtype=np.int32)),
     ],
 )
-def test_path_c_direct_rejects_partial_interval_mask(
+def test_path_c_direct_partial_interval_mask_parity(
     starts_np: np.ndarray | None,
     ends_np: np.ndarray | None,
 ) -> None:
+    """Partial (starts-only / ends-only), out-of-range, and short/empty
+    intervals all reach set-equality parity with the reference through the
+    owner-output Path C prim, including the ``-1`` sentinel for short rows."""
+
     scores = mx.array(np.array([
         [5.0, 4.0, 3.0, 2.0],
         [0.0, 1.0, 2.0, 3.0],
     ], dtype=np.float32))
     starts = None if starts_np is None else mx.array(starts_np)
     ends = None if ends_np is None else mx.array(ends_np)
-    out = mx.full((2, 2), -123, dtype=mx.int32)
 
-    with pytest.raises(TopKPathCDirectError, match="limited to unmasked rows"):
-        topk_selector_tilelang_direct(scores, k=2, starts=starts, ends=ends, out=out)
-    mx.eval(out)
-    assert np.asarray(out).tolist() == [[-123, -123], [-123, -123]]
+    out_c = _topk_tilelang_direct_output(scores, k=2, starts=starts, ends=ends)
+    out_ref = topk_selector_reference(scores, k=2, starts=starts, ends=ends)
+    mx.eval(out_c, out_ref)
+
+    assert _to_index_sets(out_c) == _to_index_sets(out_ref)
 
 
 def test_path_c_keeps_expected_exact_order_for_ties() -> None:
@@ -610,28 +616,28 @@ def test_path_c_keeps_expected_exact_order_for_ties() -> None:
     assert np.asarray(out_c).tolist() == expected
 
 
-def test_path_c_direct_rejects_masked_ties_and_sentinels() -> None:
+def test_path_c_direct_masked_ties_and_sentinels_parity() -> None:
+    """Masked rows with ties and short intervals (``-1`` sentinels) reach
+    set-equality parity with the reference through the owner-output Path C
+    prim. Tie-exact sequence ordering is not asserted (set membership is the
+    contract), but the short-interval sentinel set must contain ``-1``."""
+
     scores = mx.array(np.array([
         [9.0, 7.0, 7.0, 6.0, 7.0, 5.0, 7.0, 4.0],
         [1.0, 8.0, 8.0, 8.0, 3.0, 8.0, 2.0, 8.0],
         [0.0, 9.0, 8.0, 7.0, 6.0, 5.0, 4.0, 4.0],
     ], dtype=np.float32))
     starts = mx.array(np.array([1, 2, 6], dtype=np.int32))
     ends = mx.array(np.array([7, 6, 8], dtype=np.int32))
-    out = mx.full((3, 4), -123, dtype=mx.int32)
 
-    with pytest.raises(TopKPathCDirectError, match="limited to unmasked rows"):
-        topk_selector_tilelang_direct(scores, k=4, starts=starts, ends=ends, out=out)
+    out_c = _topk_tilelang_direct_output(scores, k=4, starts=starts, ends=ends)
     out_ref = topk_selector_reference(scores, k=4, starts=starts, ends=ends)
-    mx.eval(out, out_ref)
+    mx.eval(out_c, out_ref)
 
-    expected = [
-        [1, 2, 4, 6],
-        [2, 3, 5, 4],
-        [6, 7, -1, -1],
-    ]
-    assert np.asarray(out_ref).tolist() == expected
-    assert np.asarray(out).tolist() == [[-123] * 4] * 3
+    assert _to_index_sets(out_c) == _to_index_sets(out_ref)
+    # Row 2's interval [6, 8) has only 2 valid columns for k=4, so the
+    # reference pads with the -1 sentinel; Path C must do the same.
+    assert -1 in _to_index_sets(out_c)[2]
 
 
 @pytest.mark.parametrize("dtype", [mx.float32, mx.float16])
@@ -728,28 +734,58 @@ def fail_path_b(*args: object, **kwargs: object) -> None:
     assert _to_index_sets(out) == _to_index_sets(out_ref)
 
 
-def test_public_entry_point_auto_uses_reference_for_masked_calls(
+def test_public_entry_point_auto_uses_path_c_for_masked_calls(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None:
+    """Masked AUTO calls now route through the owner-output Path C prim
+    (passing starts/ends through), not the pure-MLX reference."""
+
     scores = mx.array(np.arange(64, dtype=np.float32).reshape(1, 64))
     starts = mx.array(np.array([8], dtype=np.int32))
     ends = mx.array(np.array([48], dtype=np.int32))
+    sentinel = mx.array(np.array([[47, 46, 45, 44, 43, 42, 41, 40]], dtype=np.int32))
+    seen: dict[str, object] = {}
 
-    def fail_path_c(*args: object, **kwargs: object) -> None:
-        raise AssertionError("masked backend='auto' must not call unmasked Path C direct")
+    def path_c_direct(scores_arg, k, *, starts=None, ends=None, out=None):
+        seen["starts"] = starts
+        seen["ends"] = ends
+        return sentinel
 
     def fail_path_b(*args: object, **kwargs: object) -> None:
         raise AssertionError("backend='auto' must not call retired Path B")
 
     monkeypatch.setattr(
         "cppmega_mlx.nn._tilelang.topk_selector.topk_selector_tilelang_direct",
-        fail_path_c,
+        path_c_direct,
     )
     monkeypatch.setattr(
         "cppmega_mlx.nn._tilelang.topk_selector.topk_selector_metal",
         fail_path_b,
     )
     out = topk_selector(scores, k=8, starts=starts, ends=ends, backend="auto")
+    assert seen["starts"] is starts
+    assert seen["ends"] is ends
+    assert np.asarray(out).tolist() == [[47, 46, 45, 44, 43, 42, 41, 40]]
+
+
+def test_public_entry_point_auto_masked_falls_back_to_reference_when_path_c_fails(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """If the masked Path C prim cannot dispatch, AUTO still falls back to the
+    pure-MLX reference rather than raising."""
+
+    scores = mx.array(np.arange(64, dtype=np.float32).reshape(1, 64))
+    starts = mx.array(np.array([8], dtype=np.int32))
+    ends = mx.array(np.array([48], dtype=np.int32))
+
+    def fail_path_c(*args: object, **kwargs: object) -> None:
+        raise TopKPathCDirectError("masked compile failed")
+
+    monkeypatch.setattr(
+        "cppmega_mlx.nn._tilelang.topk_selector.topk_selector_tilelang_direct",
+        fail_path_c,
+    )
+    out = topk_selector(scores, k=8, starts=starts, ends=ends, backend="auto")
     out_ref = topk_selector_reference(scores, k=8, starts=starts, ends=ends)
     assert _to_index_sets(out) == _to_index_sets(out_ref)