fix(path-c): split Metal FORWARD fused mega-kernel (chain_3_7) to clear MTLCompilerService XPC crash at newComputePipelineState

ROOT CAUSE (verified, full local_gb10_quarter: depth=13 hidden=3584 max_seq=4096):
The greedy direct-chain planner fuses each FORWARD segment up to the portable
kernel-buffer limit, producing the 4-op forward mega-kernel
  m2rnn + residual_rmsnorm + attention_qkv_projection + sparse_mla_fp8_apply
(region local_gb10_quarter_path_c_10_12_chain_3_7). Its generated MSL device
kernel is ~176-199 KB / ~1600-2211 lines, dominated by attention_qkv_projection
(207 refs) + sparse_mla_fp8_apply (199 refs) -- a sparse top-k FP8 attention with
RoPE + softmax inlined into ONE monolithic `kernel void`. TileLang codegen +
metallib build SUCCEED (~3s, exit 0), but at runtime Metal's MTLCompilerService
crashes the final AIR->GPU-ISA pipeline-state stage inside newComputePipelineState
(lazy, inside mx.eval at tilelang tvm_ffi adapter):
  InternalError: Check failed: (state != nullptr): ... for function
  local_gb10_quarter_path_c_10_12_chain_3_7_kernel Compilation failed due to an
  interrupted connection: XPC_ERROR_CONNECTION_INTERRUPTED.
Deterministic (3/3 runs), ~3.8s, at exec#1 on the FORWARD chain -- gates the whole
route (backward never reached). Reproduces on baseline 097693e.

Measured shader-size hierarchy (device MSL):
  chain_0_3 (3 light fwd ops)               ~46 KB / 669 L  -> pipeline OK
  chain_7_10 (3 bwd ops)                    ~116 KB / 1555 L -> pipeline OK
  chain_3_7 (4 heavy fwd ops)               ~176-199 KB     -> pipeline CRASH

FIX (option a: split the fused forward segment on Metal, explicit by target):
plan_path_c_direct_fusion_chain_for_region gains forward_max_segment_nodes, a
FORWARD-only op-count cap resolved at call time from the lowering target --
METAL_FORWARD_MAX_SEGMENT_NODES = 2 on Metal, None (uncapped) on CUDA (CUDA's
compiler has no pipeline-state size limit, so CUDA fusion is byte-for-byte
unchanged). The cap is forward-only; backward keeps greedy fusion up to the
buffer limit (its reverse-time watchdog split via launcher_chunks is separate).
Selection is explicit by target, NOT a silent runtime fallback (RULE #1).

RESULT (verified, full scale): the 4-op chain_3_7 is eliminated; the forward phase
is now 4 split segments, each well under the crashing size:
  chain_0_2 entry_rmsnorm+mamba3_mimo               ~46 KB / 709 L
  chain_2_4 residual_rmsnorm+m2rnn                   ~36 KB / 599 L
  chain_4_6 residual_rmsnorm+attention_qkv_projection ~93 KB / 861 L
  chain_6_7 sparse_mla_fp8_apply                     ~91 KB / 1056 L
All 4 forward segments now CREATE their Metal pipeline state (NO
XPC_ERROR_CONNECTION_INTERRUPTED) and EXECUTE. The full-scale forward route gets
fully past exec#1; execution proceeds into the backward phase and reaches a
DIFFERENT, separate failure -- a GPU watchdog timeout
(kIOGPUCommandBufferCallbackErrorTimeout) at the first backward segment
(chain_7_10) -- which is progress (a pre-existing backward-chain issue), not the
forward megakernel crash. The forward-only standalone executor test confirms the
forward phase executes status=ok end-to-end at smoke scale.

Side effect (benign): splitting sparse_mla_fp8_apply into its own segment makes
its int32 sparse-index scratch (path_c_int32_scratch_bank) kernel-internal rather
than a coalesced cross-segment logical scratch bank -> one fewer required logical
buffer (94->93). Updated the smoke-scale executor/route tests to the new correct
Metal grouping (7->9 segments, 2->4 forward, time-chunked bwd indices {3,5}->{5,7},
buffer counts shifted by the dropped int32 scratch bank). CUDA grouping unchanged.

Pre-existing failures NOT touched (confirmed identical on baseline 097693e):
  test_plan_path_c_direct_fusion_chains_for_model_uses_dynamic_brick_chain,
  test_direct_fusion_chain_splits_model_route_under_metal_buffer_limit,
  test_direct_fusion_chain_keeps_loss_bridge_forward_boundary_separate
  (stale "blocked" assertions -- mamba3_mimo_bwd no longer blocks),
  test_path_c_direct_chain_runtime_value_and_grad_* (broken 'tokens' batch fixture).

Author: apstenku123

cppmega_mlx/runtime/path_c_fusion_schedules.py

@@ -209,6 +209,43 @@
         DESCRIPTOR_EXECUTION_STAGE_BACKWARD,
     }
 )
+# Metal-only FORWARD fused-segment op-count cap.
+#
+# The greedy direct-chain planner fuses each forward segment up to the portable
+# kernel-buffer limit, producing the 4-op forward mega-kernel
+# ``m2rnn + residual_rmsnorm + attention_qkv_projection + sparse_mla_fp8_apply``
+# (region ``..._chain_3_7``). Its generated MSL device kernel is ~176-199 KB /
+# ~1600-2200 lines (dominated by attention_qkv_projection + sparse_mla_fp8_apply,
+# a sparse top-k FP8 attention with RoPE + softmax). TileLang codegen + metallib
+# build SUCCEED, but at runtime Metal's MTLCompilerService crashes the final
+# AIR->GPU-ISA pipeline-state stage inside ``newComputePipelineState``:
+#   InternalError: Check failed: (state != nullptr): ...
+#   Compilation failed due to an interrupted connection:
+#   XPC_ERROR_CONNECTION_INTERRUPTED. This error occurred after multiple retries.
+# Measured (full local_gb10_quarter, depth=13 hidden=3584 max_seq=4096):
+#   chain_0_3 (3 light fwd ops)  ~46 KB  -> pipeline OK
+#   chain_7_10 (3 bwd ops)       ~116 KB -> pipeline OK
+#   chain_3_7 (4 heavy fwd ops)  ~176 KB -> pipeline CRASH (newComputePipelineState)
+# Capping forward segments at 2 ops splits chain_3_7 into chain_3_5 and chain_5_7,
+# each well under the crashing size. CUDA's compiler has no such pipeline-state
+# limit, so this cap is gated to Metal only (CUDA keeps greedy forward fusion).
+METAL_FORWARD_MAX_SEGMENT_NODES = 2
+
+
+class _ResolveFromTarget:
+    """Sentinel: resolve a planner default from the active lowering target.
+
+    A distinct type (not ``None``, which is a valid "no cap" value) so callers
+    can explicitly request the target-resolved default or override it.
+    """
+
+    __slots__ = ()
+
+    def __repr__(self) -> str:  # pragma: no cover - debug aid
+        return "<resolve-from-target>"
+
+
+_RESOLVE_FORWARD_CAP_FROM_TARGET = _ResolveFromTarget()
 _GENERIC_MODEL_REAL_ABI_INPUT_SUFFIXES = (
     "residual_norm_weight",
     # Block A: the per-brick entry RMSNorm weight emitted by the model-
@@ -14190,6 +14227,9 @@ def plan_path_c_direct_fusion_chain_for_region(
     include_backward: bool = True,
     max_kernel_buffers: int = DESCRIPTOR_PORTABLE_KERNEL_BUFFER_LIMIT,
     max_segment_nodes: int | None = None,
+    forward_max_segment_nodes: int | None | _ResolveFromTarget = (
+        _RESOLVE_FORWARD_CAP_FROM_TARGET
+    ),
     registry: PathCFusionScheduleRegistry | None = None,
 ) -> PathCFusionScheduleChainPlan:
     """Greedily split a Path C region into direct-buffer fused segments.
@@ -14198,6 +14238,22 @@ def plan_path_c_direct_fusion_chain_for_region(
     would exceed Metal's portable buffer slot limit.  It never falls back to
     dtype-bank packing; segments that cannot be expressed with direct buffers
     under ``max_kernel_buffers`` are reported as blocked.
+
+    ``forward_max_segment_nodes`` is a Metal-only cap on the number of ops a
+    FORWARD-phase fused segment may contain. The default
+    (``_RESOLVE_FORWARD_CAP_FROM_TARGET``) is resolved at call time from the
+    active lowering target: on Metal it is
+    :data:`METAL_FORWARD_MAX_SEGMENT_NODES` = 2 (so the heavy 4-op forward
+    segment ``m2rnn + residual_rmsnorm + attention_qkv_projection +
+    sparse_mla_fp8_apply`` — whose ~176-199 KB MSL crashes Metal's
+    ``MTLCompilerService`` at ``newComputePipelineState`` with
+    ``XPC_ERROR_CONNECTION_INTERRUPTED`` — splits into two smaller
+    pipeline-safe kernels); on CUDA it is ``None`` (no cap, the CUDA compiler
+    has no such pipeline-state size limit, so CUDA fusion is unchanged). The
+    cap applies ONLY to forward segments; backward segments keep greedy fusion
+    up to the buffer limit (their reverse-time watchdog split is handled
+    separately via launcher_chunks). Pass an explicit value to override; pass
+    ``None`` to disable the cap entirely.
     """
 
     if not isinstance(region, PathCFusionRegion):
@@ -14206,6 +14262,23 @@ def plan_path_c_direct_fusion_chain_for_region(
         raise ValueError("max_kernel_buffers must be positive")
     if max_segment_nodes is not None and max_segment_nodes <= 0:
         raise ValueError("max_segment_nodes must be positive when provided")
+    if isinstance(forward_max_segment_nodes, _ResolveFromTarget):
+        # Resolve the forward fusion cap from the active lowering target:
+        # Metal needs the split (MTLCompilerService crashes at
+        # newComputePipelineState on the 4-op forward mega-kernel); CUDA has no
+        # such limit and stays uncapped (greedy) for max fusion.
+        forward_max_segment_nodes = (
+            METAL_FORWARD_MAX_SEGMENT_NODES
+            if _path_c_default_target() == "metal"
+            else None
+        )
+    if (
+        forward_max_segment_nodes is not None
+        and forward_max_segment_nodes <= 0
+    ):
+        raise ValueError(
+            "forward_max_segment_nodes must be positive when provided"
+        )
     working_region = (
         build_path_c_aot_autograd_region(region)
         if include_backward
@@ -14244,6 +14317,24 @@ def plan_path_c_direct_fusion_chain_for_region(
                     "execution boundary required by the loss cotangent bridge"
                 )
                 break
+            # FORWARD-only Metal pipeline-state cap. Forward fused segments
+            # larger than ``forward_max_segment_nodes`` ops generate MSL whose
+            # ~176-199 KB device kernel crashes Metal's MTLCompilerService at
+            # newComputePipelineState (XPC_ERROR_CONNECTION_INTERRUPTED). Stop
+            # extending the forward segment here so it splits into smaller
+            # pipeline-safe kernels. Backward segments are exempt (they keep
+            # greedy fusion up to the buffer limit).
+            if (
+                execution_phase == DESCRIPTOR_EXECUTION_STAGE_FORWARD
+                and forward_max_segment_nodes is not None
+                and end - start > forward_max_segment_nodes
+            ):
+                first_failure = (
+                    f"forward direct-buffer segment reached "
+                    f"forward_max_segment_nodes={forward_max_segment_nodes} "
+                    f"(Metal newComputePipelineState size cap)"
+                )
+                break
             target = selector.select(candidate_region)
             if target is None:
                 first_failure = (
@@ -14386,6 +14477,9 @@ def plan_path_c_direct_fusion_chains_for_model(
     min_route_bricks: int = 2,
     max_kernel_buffers: int = DESCRIPTOR_PORTABLE_KERNEL_BUFFER_LIMIT,
     max_segment_nodes: int | None = None,
+    forward_max_segment_nodes: int | None | _ResolveFromTarget = (
+        _RESOLVE_FORWARD_CAP_FROM_TARGET
+    ),
     registry: PathCFusionScheduleRegistry | None = None,
     sequence_length: int | None = None,
 ) -> tuple[PathCFusionScheduleChainPlan, ...]:
@@ -14395,6 +14489,10 @@ def plan_path_c_direct_fusion_chains_for_model(
     ``plan_path_c_fusion_schedules_for_model``: discover regions from the
     model's brick graph first, then split each discovered region into generic
     direct-buffer segments.  It does not consult named acceptance fixtures.
+
+    ``forward_max_segment_nodes`` is forwarded to
+    :func:`plan_path_c_direct_fusion_chain_for_region` (Metal-only forward
+    fusion cap; see that function's docstring).
     """
 
     regions = build_path_c_model_regions_from_model(
@@ -14410,6 +14508,7 @@ def plan_path_c_direct_fusion_chains_for_model(
             include_backward=include_backward,
             max_kernel_buffers=max_kernel_buffers,
             max_segment_nodes=max_segment_nodes,
+            forward_max_segment_nodes=forward_max_segment_nodes,
             registry=registry,
         )
         for region in regions

scripts/repro_fullscale_directchain.py

@@ -0,0 +1,157 @@
+#!/usr/bin/env python3
+"""DECISIVE full-scale repro: compile + execute the local_gb10_quarter Path C
+direct-chain BACKWARD segment shaders at FULL model scale (depth=13, dims=3584).
+
+This drives the SAME path as scripts/m04_train_step.py
+--use-path-c-direct-chain-runtime: it builds the full local_gb10_quarter regions,
+plans the direct fusion chain (mamba3_mimo_bwd now monolithic grid_chunks,
+m2rnn_bwd time-chunked launcher_chunks), NATIVE-COMPILES every segment shader via
+TileLang -> Metal newComputePipelineState (this is where the prior
+XPC_ERROR_CONNECTION_INTERRUPTED crash happened), then EXECUTES the route with
+caller-owned zero buffers (this is where kIOGPUCommandBufferCallbackErrorTimeout
+would fire if mamba3 monolithic trips the GPU watchdog).
+
+No kernel patching, no model code swap; routes are the production direct-chain
+schedule. RULE #1: the guarded runtime RAISES on failure (no silent fallback).
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from pathlib import Path
+import sys
+import time
+import traceback
+from typing import Any
+
+ROOT = Path(__file__).resolve().parents[1]
+if str(ROOT) not in sys.path:
+    sys.path.insert(0, str(ROOT))
+SCRIPTS = ROOT / "scripts"
+if str(SCRIPTS) not in sys.path:
+    sys.path.insert(0, str(SCRIPTS))
+
+import mlx.core as mx  # noqa: E402
+import m04_train_step as m  # noqa: E402
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--seq-len", type=int, default=None,
+                        help="sequence length for the region shape env (default: region default)")
+    parser.add_argument("--execute", action="store_true",
+                        help="after compile, EXECUTE the route (drives the GPU watchdog)")
+    args = parser.parse_args()
+
+    print(f"[repro] mx default device: {mx.default_device()}", flush=True)
+    print(f"[repro] PATH_C target: {m._path_c_default_target()}", flush=True)
+    print(f"[repro] per-segment cmdbuf commit: "
+          f"{m._path_c_per_segment_command_buffer_commit_enabled()}", flush=True)
+
+    # 1) FULL-scale region (depth=13, hidden=3584) -- the real production route.
+    profile, route_symbols, regions = m._local_gb10_path_c_model_regions()
+    print(f"[repro] profile={profile.name} depth={profile.depth} "
+          f"hidden={profile.hidden_size} max_seq={profile.max_seq_length}", flush=True)
+    sel = m._select_path_c_model_route_region(regions)
+    assert sel is not None
+    print(f"[repro] selected region={sel.name} nodes={len(sel.nodes)} "
+          f"ops={[n.op_name for n in sel.nodes]}", flush=True)
+
+    scheduled = m.plan_path_c_fusion_schedule_for_region(sel, include_backward=True)
+    chain = m.plan_path_c_direct_fusion_chain_for_region(
+        scheduled.region, include_backward=True,
+    )
+    print(f"[repro] chain.status={chain.status} n_segments={len(chain.segments)}", flush=True)
+    for seg in chain.segments:
+        tgt = getattr(seg, "schedule_target", None)
+        ops = [n.op_name for n in seg.region.nodes]
+        rd = getattr(tgt, "row_dispatch_mode", None) if tgt else None
+        print(f"[repro]   seg[{seg.index}] phase={seg.execution_phase} "
+              f"status={seg.status} row_dispatch={rd} ops={ops}", flush=True)
+
+    # 2) NATIVE COMPILE each segment shader -- the newComputePipelineState step.
+    print("\n[repro] === COMPILING segment shaders (newComputePipelineState) ===", flush=True)
+    t_compile = time.perf_counter()
+    try:
+        artifacts = m.compile_path_c_direct_fusion_chain_artifacts(chain)
+    except BaseException as exc:  # noqa: BLE001  surface the crash exactly
+        print(f"[repro] !!! COMPILE FAILED after {time.perf_counter()-t_compile:.2f}s: "
+              f"{type(exc).__name__}: {exc}", flush=True)
+        traceback.print_exc()
+        return 3
+    compile_elapsed = time.perf_counter() - t_compile
+    print(f"[repro] === ALL {len(artifacts)} segment shaders COMPILED ok in "
+          f"{compile_elapsed:.2f}s ===", flush=True)
+
+    if not args.execute:
+        print(json.dumps({
+            "phase": "compile_only",
+            "compiled_segments": len(artifacts),
+            "compile_elapsed_seconds": compile_elapsed,
+            "mamba3_seg_dispatch": next(
+                getattr(getattr(s, "schedule_target", None), "row_dispatch_mode", None)
+                for s in chain.segments
+                if any(n.op_name == "mamba3_mimo_bwd" for n in s.region.nodes)
+            ),
+            "m2rnn_seg_dispatch": next(
+                getattr(getattr(s, "schedule_target", None), "row_dispatch_mode", None)
+                for s in chain.segments
+                if any(n.op_name == "m2rnn_bwd" for n in s.region.nodes)
+            ),
+        }, indent=2))
+        return 0
+
+    # 3) EXECUTE the route -- caller-owned zero buffers; drives the GPU watchdog.
+    print("\n[repro] === EXECUTING direct-chain route (GPU watchdog window) ===", flush=True)
+    specs = m._path_c_direct_chain_required_logical_buffer_specs(chain)
+    buffers: dict[str, Any] = {}
+    for name, spec in specs.items():
+        dtype = getattr(mx, str(spec["dtype"]))
+        buffers[name] = mx.zeros(tuple(int(d) for d in spec["shape"]), dtype=dtype)
+    mx.eval(*buffers.values())
+    print(f"[repro] allocated {len(buffers)} caller-owned logical buffers", flush=True)
+
+    t_run = time.perf_counter()
+    try:
+        payload = m.run_path_c_direct_fusion_chain_route(
+            chain=chain,
+            logical_buffers=buffers,
+            artifacts=artifacts,
+        )
+    except BaseException as exc:  # noqa: BLE001
+        print(f"[repro] !!! EXECUTE FAILED after {time.perf_counter()-t_run:.2f}s: "
+              f"{type(exc).__name__}: {exc}", flush=True)
+        traceback.print_exc()
+        return 4
+    run_elapsed = time.perf_counter() - t_run
+    print(f"[repro] === ROUTE EXECUTED status={payload.get('status')} in "
+          f"{run_elapsed:.2f}s ===", flush=True)
+
+    # Per-segment timing + status (focus on mamba3_mimo_bwd / m2rnn_bwd).
+    seg_report = []
+    for s in payload.get("segments", []):
+        seg_report.append({
+            "index": s.get("index"),
+            "phase": s.get("execution_phase"),
+            "status": s.get("status"),
+            "ops": [n for n in s.get("op_signature", s.get("ops", []))]
+                   if isinstance(s.get("op_signature", s.get("ops", [])), list) else None,
+            "time_chunk_launch_count": s.get("time_chunk_launch_count"),
+            "elapsed_seconds": s.get("elapsed_seconds"),
+        })
+    print(json.dumps({
+        "phase": "executed",
+        "route_status": payload.get("status"),
+        "runtime_uses_direct_fusion_chain": payload.get("runtime_uses_direct_fusion_chain"),
+        "segment_count": payload.get("segment_count"),
+        "compile_elapsed_seconds": compile_elapsed,
+        "run_elapsed_seconds": run_elapsed,
+        "segments": seg_report,
+    }, indent=2, default=str))
+    return 0 if payload.get("status") == "ok" else 5
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())