fix(path-c): device-aware TileLang target (CUDA on non-Metal) + robust dev-build root

Path C fused compile hardcoded target=metal and the dev-build resolver pinned
TVM_LIBRARY_PATH to a stub-only dir (no libtilelang.so) on gb10, so the planner
failed and training fell back to the Metal-only eager kernel (unavailable on CUDA).
Add _path_c_default_target() (cuda when mx.metal unavailable, else metal) and use it
for all fused compile defaults; require libtilelang.so + honor TILELANG_DEV_BUILD_ROOT/
TILELANG_ROOT + derive the real tilelang build when resolving the dev env. Verified on
gb10: lowerer runs target=cuda (sm_121), fused block JIT-compiles. Apple stays metal
(no regression). Runtime install still gated by the value_and_grad grad-tree bridge wiring.

Author: apstenku123

cppmega_mlx/runtime/path_c_fusion.py

@@ -19,6 +19,19 @@
 
 import mlx.core as mx
 
+
+def _path_c_default_target() -> str:
+    """Default TileLang lowering target for Path C kernels.
+
+    Metal on Apple silicon; CUDA elsewhere (e.g. gb10).  ``mx.metal`` may exist
+    as a module yet report ``is_available() == False`` on non-Metal hosts.
+    """
+    metal = getattr(mx, "metal", None)
+    if metal is None or not metal.is_available():
+        return "cuda"
+    return "metal"
+
+
 __all__ = [
     "BenchmarkAcceptanceRow",
     "CompiledPathCRegion",
@@ -2199,7 +2212,7 @@ def compile_mamba3_fp8_train_tilelang_region_from_prim_funcs(
     fp8_prepare_spec: SparseMLAFp8PrepareSpec | None = None,
     sparse_mla_fp8_prepared: Any,
     schedule_template: Callable[[Any], Any] | None = None,
-    target: str = "metal",
+    target: str = _path_c_default_target(),
     lowerer: Callable[..., Any] | None = None,
     require_single_kernel: bool = True,
 ) -> Any:
@@ -2467,7 +2480,7 @@ def compile_path_c_region(
     schedule_status: str = "ready",
     tilelang_lowerer: Callable[..., Any] | None = None,
     tilelang_plan_factory: Callable[..., Any] | None = None,
-    target: str = "metal",
+    target: str = _path_c_default_target(),
     compiler: Callable[[FusionCompilePlan], object] | None = None,
 ) -> FusionCompilePlan | CompiledPathCRegion:
     """Create a TileLang/TVM region compile plan, optionally invoking compiler."""
@@ -2575,7 +2588,7 @@ def compile_path_c_region(
 def tilelang_single_entry_lowerer(
     func_or_mod: Any,
     *,
-    target: str = "metal",
+    target: str = _path_c_default_target(),
     execution_backend: str = "tvm_ffi",
     compile_prim_func: Callable[..., Any] | None = None,
     **_kwargs: Any,

cppmega_mlx/runtime/path_c_fusion_schedules.py

@@ -22,6 +22,7 @@
 
 from cppmega_mlx.runtime.path_c_fusion import (
     CompiledPathCRegion,
+    _path_c_default_target,
     FusionCompilePlan,
     FusionKernelSurface,
     FusionScheduleContractStatus,
@@ -13499,7 +13500,7 @@ def compile(
         self,
         *,
         tilelang_lowerer: Callable[..., Any],
-        target_name: str = "metal",
+        target_name: str = _path_c_default_target(),
     ) -> CompiledPathCRegion:
         """Compile the selected target through its descriptor schedule template."""
 
@@ -14467,7 +14468,7 @@ def plan_mamba3_fp8_train_fusion_schedule(
 def compile_mamba3_fp8_train_fusion_schedule(
     *,
     tilelang_lowerer: Callable[..., Any] | None = None,
-    target_name: str = "metal",
+    target_name: str = _path_c_default_target(),
     include_backward: bool = True,
     model_config: Any | None = None,
     max_rows_per_launch: int | None = DESCRIPTOR_DEFAULT_MAX_ROWS_PER_LAUNCH,

scripts/m04_train_step.py

@@ -1058,11 +1058,36 @@ def fp8_path_c_producer_gate_payload(
     }
 
 
+def _path_c_default_target() -> str:
+    """Resolve the default lowering target for Path C fused/direct-chain kernels.
+
+    On Apple silicon we lower to Metal; everywhere else (e.g. CUDA gb10) Metal is
+    unavailable, so fall back to CUDA.  Detection mirrors mlx: ``mx.metal`` may be
+    present as a module yet report ``is_available() == False`` on non-Metal hosts.
+    """
+    metal = getattr(mx, "metal", None)
+    if metal is None or not metal.is_available():
+        return "cuda"
+    return "metal"
+
+
 def _tilelang_dev_roots() -> tuple[Path, ...]:
     roots: list[Path] = []
-    env_root = os.environ.get("TILELANG_DEV_BUILD_ROOT")
-    if env_root:
-        roots.append(Path(env_root).expanduser())
+    for env_name in ("TILELANG_DEV_BUILD_ROOT", "TILELANG_ROOT"):
+        env_root = os.environ.get(env_name)
+        if env_root:
+            roots.append(Path(env_root).expanduser())
+    # Derive the build dir of the installed `tilelang` package (real build with
+    # libtilelang.so), e.g. /home/dave/source/tilelang/build, as a candidate.
+    try:
+        import importlib.util
+
+        spec = importlib.util.find_spec("tilelang")
+        if spec is not None and spec.origin:
+            pkg_build = Path(spec.origin).parent.parent / "build"
+            roots.append(pkg_build)
+    except Exception:
+        pass
     roots.extend(
         [
             ROOT.parent / "tl_apache_tvm_swap",
@@ -1110,6 +1135,10 @@ def ensure_tilelang_dev_env_for_path_c() -> None:
         tvm_dir = build_root / "tvm"
         if not lib_dir.exists() or not tvm_dir.exists():
             continue
+        # Skip stub-only build dirs (e.g. tl_apache_tvm_swap/build holds only
+        # libcudart_stub.so etc.). A usable dev build must ship libtilelang.so.
+        if not (lib_dir / "libtilelang.so").exists():
+            continue
         os.environ["TILELANG_DEV_BUILD_ROOT"] = str(build_root)
         os.environ.setdefault("TVM_LIBRARY_PATH", str(lib_dir))
         _prepend_env_path("DYLD_LIBRARY_PATH", lib_dir)
@@ -3225,7 +3254,7 @@ def fp8_path_c_training_route_payload_for_model(
     compile_receipt_path: str | Path | None = None,
     auto_install_fused_train_block: bool = False,
     fused_train_block_artifact_lowerer: Callable[..., Any] | None = None,
-    fused_train_block_artifact_target_name: str = "metal",
+    fused_train_block_artifact_target_name: str = _path_c_default_target(),
     fused_train_block_artifact_execution_backend: str = "tvm_ffi",
 ) -> dict[str, Any]:
     """Return Path C training route metadata using model-owned ABI banks."""
@@ -4759,7 +4788,7 @@ def compile_path_c_fused_train_block_artifact_for_model(
     *,
     model: Any,
     sequence_length: int | None = None,
-    target_name: str = "metal",
+    target_name: str = _path_c_default_target(),
     execution_backend: str = "tvm_ffi",
     lowerer: Callable[..., Any] | None = None,
 ) -> dict[str, Any]:
@@ -5556,7 +5585,7 @@ def install_path_c_fused_train_block_runtime_for_model(
     training_runtime: Any | None = None,
     compile_artifact: bool = False,
     artifact_lowerer: Callable[..., Any] | None = None,
-    artifact_target_name: str = "metal",
+    artifact_target_name: str = _path_c_default_target(),
     artifact_execution_backend: str = "tvm_ffi",
     sequence_length: int | None = None,
 ) -> dict[str, Any]:
@@ -6020,7 +6049,7 @@ def _path_c_direct_chain_artifact_for_segment(
 def compile_path_c_direct_fusion_chain_artifacts(
     chain: Any,
     *,
-    target_name: str = "metal",
+    target_name: str = _path_c_default_target(),
     execution_backend: str = "tvm_ffi",
     lowerer: Callable[..., Any] | None = None,
 ) -> tuple[Any, ...]:
@@ -6661,7 +6690,8 @@ def install_path_c_direct_chain_training_runtime_for_model(
     chain: Any | None = None,
     artifacts: Any | None = None,
     logical_owner: Any | None = None,
-    artifact_compiler: Callable[[Any], Any] = compile_path_c_direct_fusion_chain_artifacts,
+    target_name: str | None = None,
+    artifact_compiler: Callable[..., Any] | None = None,
     owner_name: str | None = None,
     sequence_length: int | None = None,
     training_critical_path: bool = False,
@@ -6733,6 +6763,13 @@ def install_path_c_direct_chain_training_runtime_for_model(
             "training_critical_path": False,
         }
 
+    resolved_target_name = target_name or _path_c_default_target()
+    if artifact_compiler is None:
+        def artifact_compiler(_chain: Any) -> Any:
+            return compile_path_c_direct_fusion_chain_artifacts(
+                _chain,
+                target_name=resolved_target_name,
+            )
     compiled_artifacts = artifacts if artifacts is not None else artifact_compiler(selected_chain)
     resolved_owner_name = owner_name or (
         f"{profile_name}.path_c_direct_fusion_chain_training_runtime"