refactor: implement continuous snake graph alignment in vbgui and enforce buffer shape safety in compilation runtime.

Author: apstenku123

cppmega_mlx/nn/_tilelang/mamba3_path_c.py

@@ -108,6 +108,10 @@
     MSLDispatchUnsupported,
     can_run_metal,
 )
+from cppmega_mlx.nn._tilelang._mlx_runtime import (
+    NativeTileLangRuntimeError,
+    _validate_owner_result,
+)
 from cppmega_mlx.nn._tilelang.mamba3 import _validate_inputs
 
 
@@ -2789,14 +2793,17 @@ def mamba3_mimo_fwd_path_c(
     if not isinstance(out_list, (list, tuple)) or len(out_list) != 2:
         raise RuntimeError("Mamba3 Path C fwd tvm-ffi returned an invalid output tuple")
     if owner_outputs is not None:
-        y, h_last = owner_outputs
-        if not all(
-            got is expected
-            for got, expected in zip(out_list, (y, h_last), strict=True)
-        ):
+        try:
+            y, h_last = cast(
+                tuple[mx.array, mx.array],
+                _validate_owner_result(out_list, owner_outputs),
+            )
+        except NativeTileLangRuntimeError as exc:
             raise RuntimeError(
                 "Mamba3 Path C fwd tvm-ffi did not return caller-owned outputs"
-            )
+            ) from exc
+        del lowering
+        return y, h_last
     y, h_last = cast(tuple[mx.array, mx.array], tuple(out_list))
     del lowering
     return y, h_last

cppmega_mlx/runtime/path_c_fusion_launcher.py

@@ -114,6 +114,8 @@ class PathCAbiManifest:
     row_subchunk_index_param: str | None
     row_chunk_index_buffer: str | None
     backward_gate_param: str | None
+    backward_stage_count: int | None
+    backward_stage_index_param: str | None
 
 
 def _optional_positive_int_attr(prim_func: Any, key: str) -> int | None:
@@ -241,6 +243,12 @@ def load_path_c_abi_manifest(prim_func: Any) -> PathCAbiManifest:
         backward_gate_param=_optional_string_attr(
             prim_func, "tl.fusion.backward_gate_param"
         ),
+        backward_stage_count=_optional_positive_int_attr(
+            prim_func, "tl.fusion.backward_stage_count"
+        ),
+        backward_stage_index_param=_optional_string_attr(
+            prim_func, "tl.fusion.backward_stage_index_param"
+        ),
     )
 
 
@@ -534,6 +542,7 @@ def __call__(
         row_chunk_param_index: int | None = None
         row_chunk_param_is_scalar = False
         row_subchunk_param_index: int | None = None
+        backward_stage_param_index: int | None = None
         backward_gate_param_index: int | None = None
         for param_name in manifest.param_order:
             logical_name = param_name[: -len("_handle")] if param_name.endswith("_handle") else param_name
@@ -554,6 +563,9 @@ def __call__(
             elif logical_name == manifest.row_subchunk_index_param:
                 row_subchunk_param_index = len(positional)
                 positional.append(0)
+            elif logical_name == manifest.backward_stage_index_param:
+                backward_stage_param_index = len(positional)
+                positional.append(0)
             elif logical_name == manifest.backward_gate_param:
                 backward_gate_param_index = len(positional)
                 positional.append(_BACKWARD_GATE if run_backward else _FORWARD_ONLY_GATE)
@@ -608,6 +620,8 @@ def apply_returned_outputs(returned: Any) -> None:
                     continue
                 elif logical_name == manifest.row_subchunk_index_param:
                     continue
+                elif logical_name == manifest.backward_stage_index_param:
+                    continue
                 elif logical_name == manifest.backward_gate_param:
                     continue
                 positional[param_index] = value
@@ -678,26 +692,52 @@ def restore_logical_buffers(snapshots: Mapping[str, mx.array]) -> None:
                     "compiled PrimFunc declares row subchunk dispatch metadata "
                     "but has no row subchunk index parameter"
                 )
-            def launch_chunk(chunk_index: int, subchunk_index: int) -> None:
+            if (
+                manifest.backward_stage_index_param is not None
+                and backward_stage_param_index is None
+            ):
+                raise ValueError(
+                    "compiled PrimFunc declares backward stage dispatch metadata "
+                    "but has no backward stage index parameter"
+                )
+
+            def launch_chunk(
+                chunk_index: int,
+                subchunk_index: int,
+                backward_stage_index: int = 0,
+            ) -> None:
                 positional[row_chunk_param_index] = (
                     int(chunk_index)
                     if row_chunk_param_is_scalar
                     else mx.array([chunk_index], dtype=mx.int32)
                 )
                 if row_subchunk_param_index is not None:
                     positional[row_subchunk_param_index] = int(subchunk_index)
+                if backward_stage_param_index is not None:
+                    positional[backward_stage_param_index] = int(backward_stage_index)
                 apply_returned_outputs(self._kernel(*positional))
 
             if run_backward and backward_gate_param_index is not None:
                 positional[backward_gate_param_index] = _BACKWARD_FORWARD_PREPASS_GATE
+                prepass_subchunk_count = int(manifest.row_subchunk_count or 1)
                 for chunk_index in range(chunk_count):
-                    launch_chunk(chunk_index, 0)
+                    for subchunk_index in range(prepass_subchunk_count):
+                        launch_chunk(chunk_index, subchunk_index)
                 final_forward_states = snapshot_logical_buffers(
                     _STATE_AFTER_LOGICAL_BUFFERS
                 )
                 restore_logical_buffers(initial_backward_states)
                 positional[backward_gate_param_index] = _BACKWARD_GATE
-                launch_chunk(0, 0)
+                backward_subchunk_count = int(manifest.row_subchunk_count or 1)
+                backward_stage_count = int(manifest.backward_stage_count or 1)
+                for backward_stage_index in range(backward_stage_count):
+                    for chunk_index in range(chunk_count):
+                        for subchunk_index in range(backward_subchunk_count):
+                            launch_chunk(
+                                chunk_index,
+                                subchunk_index,
+                                backward_stage_index,
+                            )
                 restore_logical_buffers(final_forward_states)
             else:
                 for chunk_index in range(chunk_count):

cppmega_mlx/runtime/path_c_fusion_schedules.py

@@ -113,7 +113,7 @@
 PATH_C_DESCRIPTOR_SCHEDULE_GENERATOR = "dynamic_brick_descriptor_generator"
 DESCRIPTOR_DEFAULT_BUFFER_EXTENT = 4
 DESCRIPTOR_DEFAULT_THREADS = 256
-DESCRIPTOR_ROW_PHASED_THREADS = 128
+DESCRIPTOR_ROW_PHASED_THREADS = 1024
 DESCRIPTOR_INTERNAL_BUFFER_POLICY_SCALAR_LOCAL = "scalar_local"
 DESCRIPTOR_INTERNAL_BUFFER_POLICY_ROW_LOCAL_HIDDEN = "row_local_hidden"
 DESCRIPTOR_LOOP_POLICY_FLAT = "flat"
@@ -123,7 +123,7 @@
 DESCRIPTOR_PHYSICAL_ABI_POLICY_BANKED_BY_ROLE = "banked_by_role"
 DESCRIPTOR_PORTABLE_KERNEL_BUFFER_LIMIT = 31
 DESCRIPTOR_SHARED_SCRATCH_SPILL_THRESHOLD_BYTES = 1024
-DESCRIPTOR_SHARED_SCRATCH_BUDGET_BYTES = 12 * 1024
+DESCRIPTOR_SHARED_SCRATCH_BUDGET_BYTES = 8 * 1024
 MAMBA3_BWD_REPLAY_CHECKPOINT_INTERVAL = 8
 M2RNN_BWD_REPLAY_CHECKPOINT_INTERVAL = 1
 _DESCRIPTOR_ROOT_READS_MARKER = "cppmega_path_c_root_reads"
@@ -1780,6 +1780,7 @@ def build_path_c_descriptor_prim_func(
         internal_buffers=internal_buffers,
         internal_buffer_shapes=internal_buffer_shapes,
         physical_abi_plan=physical_abi_plan,
+        physical_abi_policy=validated_physical_abi_policy,
         internal_buffer_policy=validated_internal_buffer_policy,
         loop_policy=validated_loop_policy,
         max_rows_per_launch=validated_max_rows_per_launch,
@@ -2408,6 +2409,7 @@ def _descriptor_prim_func_source(
     internal_buffers: Sequence[str],
     internal_buffer_shapes: Mapping[str, tuple[int, ...]],
     physical_abi_plan: _PhysicalAbiPlan,
+    physical_abi_policy: str,
     internal_buffer_policy: str,
     loop_policy: str,
     max_rows_per_launch: int | None,
@@ -2667,6 +2669,9 @@ def _descriptor_prim_func_source(
         existing_parameter_count=len(param_lines),
         internal_buffer_names=frozenset(active_internal_buffers),
         force_spill_names=frozenset(force_spilled_internal_buffers),
+        force_builtin_spill_names=(
+            physical_abi_policy != DESCRIPTOR_PHYSICAL_ABI_POLICY_DIRECT
+        ),
     )
 
 
@@ -3861,6 +3866,7 @@ def _spill_large_shared_scratch_to_abi(
     existing_parameter_count: int,
     internal_buffer_names: frozenset[str] = frozenset(),
     force_spill_names: frozenset[str] = frozenset(),
+    force_builtin_spill_names: bool = True,
 ) -> tuple[str, Mapping[str, Mapping[str, Any]]]:
     candidates: list[dict[str, Any]] = []
     total_shared_bytes = 0
@@ -3881,7 +3887,10 @@ def _spill_large_shared_scratch_to_abi(
         force_scratch_abi = (
             scratch_name in force_spill_names
             or _is_row_phased_bwd_scratch_abi_buffer(scratch_name)
-            or _force_spill_shared_scratch_name(scratch_name)
+            or (
+                force_builtin_spill_names
+                and _force_spill_shared_scratch_name(scratch_name)
+            )
         )
         if (
             byte_count >= DESCRIPTOR_SHARED_SCRATCH_SPILL_THRESHOLD_BYTES
@@ -6363,6 +6372,7 @@ def _append_row_phased_mamba3_body(
             f"({out_inner}[{feature}] + ({d_skip} * {x_val})) * "
             f"{z_val} * (1.0 / (1.0 + T.exp(-{z_val})))"
         )
+        body.append(f"{indent * 3}T.sync_threads()")
         body.append(f"{indent * 3}if ((row + 1) % {checkpoint_interval}) == 0:")
         body.append(f"{indent * 4}{checkpoint_idx} = (row + 1) // {checkpoint_interval}")
         body.append(
@@ -9447,14 +9457,17 @@ def _append_row_phased_m2rnn_bwd_body(
         )
     body.append(f"{indent * 5}for {state_idx} in T.serial(0, {full_state_extent}):")
     body.append(f"{indent * 6}{dh}[{state_idx}] = 0.0")
-    body.append(f"{indent * 3}if lane == 0:")
-    body.append(f"{indent * 4}for {time_rev} in T.serial({time_rev_range}):")
-    body.append(f"{indent * 6}{time_idx} = {sequence_length - 1} - {time_rev}")
     if hidden_grad is not None:
-        body.append(f"{indent * 6}for {hidden_loop} in T.serial(0, {hidden_size}):")
         body.append(
-            f"{indent * 7}{_indexed_buffer_ref(hidden_grad, access_by_buffer, f'{time_idx} * {hidden_size} + {hidden_loop}')} = 0.0"
+            f"{indent * 5}for {grad_flat} in T.serial(0, "
+            f"{sequence_length * hidden_size}):"
         )
+        body.append(
+            f"{indent * 6}{_indexed_buffer_ref(hidden_grad, access_by_buffer, grad_flat)} = 0.0"
+        )
+    body.append(f"{indent * 3}if lane == 0:")
+    body.append(f"{indent * 4}for {time_rev} in T.serial({time_rev_range}):")
+    body.append(f"{indent * 6}{time_idx} = {sequence_length - 1} - {time_rev}")
     body.append(f"{indent * 6}{checkpoint_idx} = {time_idx} // {checkpoint_interval}")
     body.append(f"{indent * 6}{checkpoint_start} = {checkpoint_idx} * {checkpoint_interval}")
     body.append(f"{indent * 6}for {state_idx} in T.serial(0, {full_state_extent}):")
@@ -9734,10 +9747,7 @@ def _append_row_phased_m2rnn_bwd_body(
     body.append(f"{indent * 6}for {conv_ch} in T.serial(0, {conv_dim}):")
     body.append(f"{indent * 7}{conv_grad}[{conv_ch}] = 0.0")
     body.append(f"{indent * 6}T.sync_threads()")
-    body.append(
-        f"{indent * 6}for {proj_dim} in T.serial(lane, {in_proj_dim}, "
-        f"step={thread_count}):"
-    )
+    body.append(f"{indent * 6}for {proj_dim} in T.serial(0, {in_proj_dim}):")
     body.append(f"{indent * 7}{project_grad}[{proj_dim}] = 0.0")
     body.append(f"{indent * 6}for {feature} in T.serial(0, {features}):")
     body.append(f"{indent * 7}{head} = {feature} // {v_dim}")
@@ -9838,13 +9848,7 @@ def _append_row_phased_m2rnn_bwd_body(
         f"({h_prev}[{head} * {k_dim * v_dim} + {kk} * {v_dim} + {vv_inner}] * "
         f"{state_weight_expr_inner})"
     )
-    body.append(
-        f"{indent * 9}{tanh_val}[0] = "
-        f"({h_next}[{head} * {k_dim * v_dim} + {kk} * {v_dim} + {vv}] - "
-        f"({decay}[0] * "
-        f"{h_prev}[{head} * {k_dim * v_dim} + {kk} * {v_dim} + {vv}])) / "
-        f"(1.0 - {decay}[0])"
-    )
+    body.append(f"{indent * 9}{tanh_val}[0] = T.tanh({accum}[0] + ({k_val} * {v_val}))")
     body.append(
         f"{indent * 9}{scalar0}[0] = {scalar0}[0] + "
         f"({dh}[{head} * {k_dim * v_dim} + {kk} * {v_dim} + {vv}] * "
@@ -10048,10 +10052,7 @@ def _append_row_phased_m2rnn_bwd_body(
         f"{project_grad}[{conv_ch}] + ({scalar1}[0] * {current_conv_weight_expr})"
     )
     body.append(f"{indent * 6}T.sync_threads()")
-    body.append(
-        f"{indent * 6}for {proj_dim} in T.serial(lane, {in_proj_dim}, "
-        f"step={thread_count}):"
-    )
+    body.append(f"{indent * 6}for {proj_dim} in T.serial(0, {in_proj_dim}):")
     body.append(f"{indent * 7}for {hidden_loop} in T.serial(0, {hidden_size}):")
     hidden_expr = _node_indexed_canonical_or_positional_input_expr(
         node,
@@ -10766,7 +10767,6 @@ def _mamba3_emit_recompute_row(
         )
     if (
         hidden_grad is not None
-        and not launcher_chunked_rows
         and not _is_full_sequence_bank_slot(
         hidden_grad,
         access_by_buffer,
@@ -10819,14 +10819,6 @@ def _mamba3_emit_recompute_row(
     body.append(f"{indent * 5}T.sync_threads()")
     body.append(f"{indent * 5}for {time_rev} in T.serial({time_rev_range}):")
     body.append(f"{indent * 6}{time_idx} = {sequence_length - 1} - {time_rev}")
-    if launcher_chunked_rows and hidden_grad is not None:
-        body.append(
-            f"{indent * 6}for {hidden_loop} in T.serial(lane, {hidden_size}, "
-            f"step={thread_count}):"
-        )
-        body.append(
-            f"{indent * 7}{_indexed_buffer_ref(hidden_grad, access_by_buffer, f'{time_idx} * {hidden_size} + {hidden_loop}')} = 0.0"
-        )
     body.append(f"{indent * 6}{checkpoint_idx} = {time_idx} // {checkpoint_interval}")
     body.append(
         f"{indent * 6}{checkpoint_start} = {checkpoint_idx} * {checkpoint_interval}"
@@ -13728,6 +13720,12 @@ def _physical_abi_policy_for_region(
         external_buffer_count += len(_TRAIN_STEP_SCALAR_OUTPUT_ABI_NAMES)
         external_buffer_count += len(_TRAIN_STEP_SUFFIX_LOSS_INPUT_ABI_NAMES)
         external_buffer_count += len(_TRAIN_STEP_SUFFIX_LOSS_PARAMETER_GRAD_ABI_NAMES)
+    if (
+        shape_env is not None
+        and internal_buffer_policy == DESCRIPTOR_INTERNAL_BUFFER_POLICY_ROW_LOCAL_HIDDEN
+        and loop_policy == DESCRIPTOR_LOOP_POLICY_ROW_PHASED_HIDDEN
+    ):
+        return DESCRIPTOR_PHYSICAL_ABI_POLICY_BANKED_BY_ROLE
     if external_buffer_count > DESCRIPTOR_PORTABLE_KERNEL_BUFFER_LIMIT:
         return DESCRIPTOR_PHYSICAL_ABI_POLICY_BANKED_BY_ROLE
     return DESCRIPTOR_PHYSICAL_ABI_POLICY_DIRECT

cppmega_mlx/training/compiled.py

@@ -332,7 +332,12 @@ def _kernel_exact_buffers(self, buffers: Mapping[str, Any]) -> Mapping[str, Any]
                 expected_size *= int(dim)
             actual_size = int(getattr(value, "size", 0) or 0)
             if actual_size < expected_size:
-                continue
+                raise ValueError(
+                    "physical ABI runtime bindings are not executable: "
+                    f"{name}: caller-owned kernel buffer shape {actual_shape} "
+                    f"has {actual_size} elements, expected at least "
+                    f"{expected_shape} ({expected_size} elements)"
+                )
             view = mx.reshape(value, (-1,))[:expected_size]
             if expected_shape != (expected_size,):
                 view = mx.reshape(view, expected_shape)

cppmega_v4/_tilelang/_path_d_deps.py

@@ -361,7 +361,7 @@ def path_d_imports_allowed(mode: str = "frontend") -> tuple[bool, str]:
 def unsafe_triton_frontend_import_disabled_reason(root: str | None) -> str:
     root_text = f" root={root}" if root else " root=<not found>"
     return (
-        "unsafe triton frontend import disabled because "
+        "unsafe triton frontend import disabled (not importable by default) because "
         f"{TRITON_FRONTEND_UNSAFE_IMPORT_ENV}=1 is not set; native import "
         "preflight did not authorize in-process import; Path D runtime "
         f"adapter not reached;{root_text}"
@@ -371,7 +371,7 @@ def unsafe_triton_frontend_import_disabled_reason(root: str | None) -> str:
 def unsafe_fla_import_disabled_reason(root: str | None) -> str:
     root_text = f" root={root}" if root else " root=<not found>"
     return (
-        "unsafe FLA import disabled because "
+        "unsafe FLA import disabled (not importable by default) because "
         f"{TRITON_FRONTEND_UNSAFE_IMPORT_ENV}=1 is not set; native import "
         "preflight did not authorize in-process import; Path D runtime "
         f"adapter not reached;{root_text}"

cppmega_v4/fusion/auto_planner.py

@@ -6,8 +6,8 @@
 
   - Eligibility check delegates entirely to
     :func:`cppmega_v4.fusion.compatibility.can_fuse_pair` (table-driven).
-  - A region is **extended** when the trailing brick can fuse with the
-    next candidate AND no hard limit is breached.
+  - A region is **extended** when every existing brick in the open region can
+    fuse with the next candidate AND no hard limit is breached.
   - A region is **closed** (and a new one started at the next brick)
     when eligibility says ``can_fuse=False``, when adding the next brick
     would exceed ``max_region_size`` bricks, or when the estimated
@@ -197,7 +197,9 @@ def plan_fusion_regions(
 
     The decision is:
 
-      1. If ``can_fuse_pair(tail, next).can_fuse`` is False → close.
+      1. If any existing brick in the open region cannot fuse with ``next``
+         → close.  This prevents light ops such as RMSNorm from bridging two
+         hard-incompatible heavy blocks into one invalid region.
       2. If extending would put the region above ``max_region_size`` → close.
       3. If extending would push estimated threadgroup memory above
          ``max_shared_mem_bytes`` → close.
@@ -234,12 +236,33 @@ def plan_fusion_regions(
         if not elig.can_fuse:
             close = True
             reason = elig.reason
-        elif len(open_region.nodes) + 1 > max_region_size:
+        else:
+            for existing in open_region.nodes[:-1]:
+                region_elig = can_fuse_pair(existing, current)
+                if not region_elig.can_fuse:
+                    close = True
+                    reason = (
+                        "region incompatibility: "
+                        f"{existing.kind}->{current.kind}: {region_elig.reason}"
+                    )
+                    break
+                if region_elig.backend != elig.backend:
+                    close = True
+                    reason = (
+                        "region backend mismatch: "
+                        f"{existing.kind}->{current.kind} uses "
+                        f"{region_elig.backend!r}, tail pair uses {elig.backend!r}"
+                    )
+                    break
+        if not close and len(open_region.nodes) + 1 > max_region_size:
             close = True
             reason = (
                 f"region would exceed max_region_size={max_region_size}"
             )
-        elif open_region.shared_mem_bytes + next_shared > max_shared_mem_bytes:
+        elif (
+            not close
+            and open_region.shared_mem_bytes + next_shared > max_shared_mem_bytes
+        ):
             close = True
             reason = (
                 "region would exceed shared-mem budget "