feat: implement fused train-block runtime and coalesced scratch bank allocation for FP8 path-c training

Author: apstenku123

cppmega_mlx/nn/m2rnn.py

@@ -12,6 +12,7 @@
 from cppmega_mlx.nn.mamba3 import causal_depthwise_conv1d
 
 DEFAULT_CHUNK_SIZE = 128
+PATH_B_REFERENCE_BASELINE_KERNEL = "reference_path_b_baseline"
 
 
 def _require_positive_int(name: str, value: int) -> None:
@@ -309,8 +310,9 @@ def _dispatch_m2rnn_scan(
 
     AUTO uses the retired Path B status seam only when it is available,
     otherwise it falls back to the pure-MLX :func:`chunked_m2rnn_scan`.
-    PATH_B and PATH_C are explicit and fail-closed when their native routes
-    are unavailable.
+    PATH_B keeps the benchmark baseline runnable by using the legacy
+    differentiable compatibility wrapper when direct MSL is retired. PATH_C
+    remains explicit and fail-closed when its native route is unavailable.
     """
 
     from cppmega_mlx.runtime.kernel_policy import (
@@ -369,11 +371,7 @@ def _dispatch_m2rnn_scan(
     )
 
     status = m2rnn_metal_status(q)
-    if path is KernelPath.PATH_B and not status.available:
-        raise RuntimeError(
-            f"m2rnn: Path B kernel unavailable ({status.reason})"
-        )
-    if status.available:
+    if path is KernelPath.PATH_B or status.available:
         q_b, k_b, v_b, W_b, xf_b = broadcast_m2rnn_heads(q, k, v, W, xf)
         batch, _seq, heads, k_dim = q_b.shape
         v_dim = v_b.shape[-1]
@@ -386,7 +384,15 @@ def _dispatch_m2rnn_scan(
             dtype=q_b.dtype,
         )
         out, h = m2rnn_apply_with_state(q_b, k_b, v_b, W_b, xf_b, h0_full)
-        record_dispatch("m2rnn", path, "metal_kernel_fwd_v1")
+        record_dispatch(
+            "m2rnn",
+            path,
+            (
+                "metal_kernel_fwd_v1"
+                if status.available
+                else PATH_B_REFERENCE_BASELINE_KERNEL
+            ),
+        )
         return out, h
 
     record_dispatch("m2rnn", path, "reference_pure_mlx")
@@ -782,6 +788,7 @@ def __call__(
     "M2RNNConfig",
     "M2RNNMixer",
     "M2RNNMixerState",
+    "PATH_B_REFERENCE_BASELINE_KERNEL",
     "broadcast_m2rnn_heads",
     "chunked_m2rnn_scan",
     "m2rnn_softplus_decay_gate",

cppmega_mlx/runtime/path_c_fusion_schedules.py

@@ -112,8 +112,8 @@
 DESCRIPTOR_PHYSICAL_ABI_POLICY_BANKED_BY_DTYPE = "banked_by_dtype"
 DESCRIPTOR_PHYSICAL_ABI_POLICY_BANKED_BY_ROLE = "banked_by_role"
 DESCRIPTOR_PORTABLE_KERNEL_BUFFER_LIMIT = 31
-DESCRIPTOR_SHARED_SCRATCH_SPILL_THRESHOLD_BYTES = 4 * 1024
-DESCRIPTOR_SHARED_SCRATCH_BUDGET_BYTES = 24 * 1024
+DESCRIPTOR_SHARED_SCRATCH_SPILL_THRESHOLD_BYTES = 1024
+DESCRIPTOR_SHARED_SCRATCH_BUDGET_BYTES = 12 * 1024
 MAMBA3_BWD_REPLAY_CHECKPOINT_INTERVAL = 64
 _DESCRIPTOR_ROOT_READS_MARKER = "cppmega_path_c_root_reads"
 _DESCRIPTOR_ROOT_WRITES_MARKER = "cppmega_path_c_root_writes"
@@ -139,6 +139,20 @@
         "mamba3_delta_grad",
     }
 )
+DESCRIPTOR_ROW_PHASED_BWD_SCRATCH_ABI_CANONICALS = frozenset(
+    {
+        "hidden",
+        "attention_hidden",
+        "hidden_after_mamba3",
+        "m2rnn_hidden",
+        "m2rnn_delta",
+        "mamba3_delta",
+        "q_fp8",
+        "q_scale",
+        "kv_fp8",
+        "kv_scale",
+    }
+)
 DESCRIPTOR_DEFAULT_MAX_ROWS_PER_LAUNCH = 64
 DESCRIPTOR_ROW_DISPATCH_GRID_CHUNKS = "grid_chunks"
 DESCRIPTOR_ROW_DISPATCH_LAUNCHER_CHUNKS = "launcher_chunks"
@@ -3033,9 +3047,9 @@ def _spill_large_shared_scratch_to_abi(
         byte_count = _flattened_extent(shape) * _DTYPE_NBYTES[dtype]
         total_shared_bytes += byte_count
         scratch_name = match.group("name")
-        force_scratch_abi = scratch_name in DESCRIPTOR_ROW_PHASED_BWD_SCRATCH_ABI_BUFFERS
+        force_scratch_abi = _is_row_phased_bwd_scratch_abi_buffer(scratch_name)
         if (
-            byte_count > DESCRIPTOR_SHARED_SCRATCH_SPILL_THRESHOLD_BYTES
+            byte_count >= DESCRIPTOR_SHARED_SCRATCH_SPILL_THRESHOLD_BYTES
             or force_scratch_abi
         ):
             candidates.append(
@@ -3178,6 +3192,21 @@ def rewrite_scratch_refs(line: str) -> str:
     return "\n".join(rewritten) + "\n", spilled
 
 
+def _is_row_phased_bwd_scratch_abi_buffer(buffer_name: str) -> bool:
+    """Return True for internal bwd scratch that must be a device ABI buffer."""
+
+    name = str(buffer_name)
+    if name in DESCRIPTOR_ROW_PHASED_BWD_SCRATCH_ABI_BUFFERS:
+        return True
+    if not name.endswith("_grad"):
+        return False
+    canonical = _canonical_buffer_name(name)
+    if canonical in DESCRIPTOR_ROW_PHASED_BWD_SCRATCH_ABI_CANONICALS:
+        return True
+    base = name[: -len("_grad")]
+    return base.endswith(("_hidden", "_delta", "_out"))
+
+
 def _validated_internal_buffer_policy(policy: str) -> str:
     normalized = str(policy)
     if normalized not in {
@@ -3594,6 +3623,27 @@ def _row_local_internal_buffer_shape(
     shape_env: PathCModelShapeEnv,
 ) -> tuple[int, ...]:
     canonical_name = _canonical_buffer_name(buffer_name)
+    if str(buffer_name).endswith("_grad") and (
+        canonical_name in DESCRIPTOR_ROW_PHASED_BWD_SCRATCH_ABI_CANONICALS
+        or str(buffer_name)[: -len("_grad")].endswith(("_hidden", "_delta", "_out"))
+    ):
+        if canonical_name == "q_fp8":
+            return (
+                shape_env.sequence_length
+                * shape_env.attention_num_q_heads
+                * shape_env.attention_head_dim,
+            )
+        if canonical_name == "q_scale":
+            return (shape_env.sequence_length * shape_env.attention_num_q_heads,)
+        if canonical_name == "kv_fp8":
+            return (
+                shape_env.sequence_length
+                * shape_env.attention_num_kv_heads
+                * shape_env.attention_head_dim,
+            )
+        if canonical_name == "kv_scale":
+            return (shape_env.sequence_length * shape_env.attention_num_kv_heads,)
+        return (shape_env.sequence_length * shape_env.hidden_size,)
     if canonical_name in {
         "mamba3_delta",
         "m2rnn_hidden",
@@ -3616,23 +3666,6 @@ def _row_local_internal_buffer_shape(
             * shape_env.attention_num_kv_heads
             * shape_env.attention_sparse_topk,
         )
-    if str(buffer_name).endswith("_grad"):
-        if canonical_name == "q_fp8":
-            return (
-                shape_env.sequence_length
-                * shape_env.attention_num_q_heads
-                * shape_env.attention_head_dim,
-            )
-        if canonical_name == "q_scale":
-            return (shape_env.sequence_length * shape_env.attention_num_q_heads,)
-        if canonical_name == "kv_fp8":
-            return (
-                shape_env.sequence_length
-                * shape_env.attention_num_kv_heads
-                * shape_env.attention_head_dim,
-            )
-        if canonical_name == "kv_scale":
-            return (shape_env.sequence_length * shape_env.attention_num_kv_heads,)
     if canonical_name == "kv_fp8":
         return (shape_env.attention_num_kv_heads * shape_env.attention_head_dim,)
     if canonical_name == "kv_scale":
@@ -12290,13 +12323,14 @@ def _dynamic_descriptor_target_for_region(
         acceptance_profile,
         descriptors,
     )
+    train_step_output_abi = _region_requires_train_step_output_abi(region, nodes)
     physical_abi_policy = _physical_abi_policy_for_region(
         nodes,
         shape_env=shape_env,
         internal_buffer_policy=internal_buffer_policy,
         loop_policy=loop_policy,
+        train_step_output_abi=train_step_output_abi,
     )
-    train_step_output_abi = _region_requires_train_step_output_abi(region, nodes)
     schedule_name = (
         acceptance_profile.schedule_name
         if acceptance_profile is not None
@@ -12412,6 +12446,7 @@ def _physical_abi_policy_for_region(
     shape_env: PathCModelShapeEnv | None,
     internal_buffer_policy: str,
     loop_policy: str,
+    train_step_output_abi: bool = False,
 ) -> str:
     internal_buffers = _internal_buffers_for_nodes(
         nodes,
@@ -12420,6 +12455,10 @@ def _physical_abi_policy_for_region(
         loop_policy=loop_policy,
     )
     external_buffer_count = len(_external_buffers_for_nodes(nodes, internal_buffers))
+    if train_step_output_abi:
+        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 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

@@ -21,7 +21,6 @@
 
 from cppmega_mlx.data.batch import LMTokenBatch, ensure_lm_batch
 from cppmega_mlx.runtime.path_c_physical_abi import (
-    physical_abi_full_runtime_kernel_args,
     physical_abi_runtime_kernel_args,
 )
 from cppmega_mlx.training.loss import next_token_cross_entropy
@@ -53,6 +52,11 @@
 PATH_C_FUSED_TRAIN_BLOCK_VALUE_AND_GRAD_CONTRACT = (
     "path_c_fused_train_block_value_and_grad_v1"
 )
+PATH_C_SCALAR_KERNEL_PARAM_DEFAULTS: Mapping[str, int] = {
+    "path_c_run_backward": 1,
+    "path_c_row_chunk_index": 0,
+    "path_c_row_subchunk_index": 0,
+}
 
 REGIONAL_COMPILE_TARGETS: Mapping[CompileTarget, bool] = {
     "mamba3_pre": True,
@@ -213,12 +217,35 @@ def _bank_buffers(self, bank_owner: Any) -> Mapping[str, Any]:
 
     def _kernel_args_from_bank_owner(self, bank_owner: Any) -> tuple[Any, ...]:
         if self.kernel_buffer_order is not None:
-            return physical_abi_full_runtime_kernel_args(
+            buffers = self._bank_buffers(bank_owner)
+            physical_abi_runtime_kernel_args(
                 self.physical_abi_map,
                 self.physical_abi_shapes,
-                self.kernel_buffer_order,
-                self._bank_buffers(bank_owner),
+                {
+                    name: buffers[name]
+                    for name in self.physical_abi_shapes
+                    if name in buffers
+                },
             )
+            missing: list[str] = []
+            args: list[Any] = []
+            for name in self.kernel_buffer_order:
+                if name in buffers:
+                    args.append(buffers[name])
+                    continue
+                if name in PATH_C_SCALAR_KERNEL_PARAM_DEFAULTS:
+                    args.append(PATH_C_SCALAR_KERNEL_PARAM_DEFAULTS[name])
+                    continue
+                missing.append(name)
+            if missing:
+                raise ValueError(
+                    "physical ABI runtime bindings are not executable: "
+                    + "; ".join(
+                        f"{name}: missing caller-owned kernel buffer"
+                        for name in missing
+                    )
+                )
+            return tuple(args)
         return physical_abi_runtime_kernel_args(
             self.physical_abi_map,
             self.physical_abi_shapes,

cppmega_mlx/training/path_c_fused_suffix.py

@@ -40,6 +40,21 @@
 
 
 FusedSuffixCallable = Callable[..., Any]
+_MX_DTYPE_BY_ABI_NAME = {
+    "bool": mx.bool_,
+    "uint8": mx.uint8,
+    "int8": mx.int8,
+    "float16": mx.float16,
+    "bfloat16": mx.bfloat16,
+    "uint16": mx.uint16,
+    "int16": mx.int16,
+    "float32": mx.float32,
+    "uint32": mx.uint32,
+    "int32": mx.int32,
+    "float64": mx.float64,
+    "uint64": mx.uint64,
+    "int64": mx.int64,
+}
 
 
 def _zeros_like_with_cotangent_dtype(value: mx.array) -> mx.array:
@@ -56,6 +71,22 @@ def _zeros_like_with_cotangent_dtype(value: mx.array) -> mx.array:
     return mx.zeros_like(value)
 
 
+def _target_mask_for_abi(
+    abi_map: Mapping[str, Any],
+    logical_name: str,
+    value: mx.array,
+) -> mx.array:
+    info = abi_map.get(logical_name)
+    if not isinstance(info, Mapping):
+        return value
+    expected_dtype = _MX_DTYPE_BY_ABI_NAME.get(str(info.get("dtype", "")))
+    if expected_dtype is None or value.dtype == expected_dtype:
+        return value
+    if value.dtype not in (mx.float16, mx.bfloat16, mx.float32, mx.float64):
+        return value
+    return value.astype(expected_dtype)
+
+
 def build_fused_suffix_custom_function(
     *,
     artifact: Any,
@@ -115,6 +146,11 @@ def _write_runtime_state(
         write_into_bank_slot(
             abi_map, bank_buffers, target_ids_logical_name, target_ids
         )
+        target_mask = _target_mask_for_abi(
+            abi_map,
+            target_mask_logical_name,
+            target_mask,
+        )
         write_into_bank_slot(
             abi_map, bank_buffers, target_mask_logical_name, target_mask
         )

cppmega_v4/_tilelang/_path_d_deps.py

@@ -13,6 +13,7 @@
 
 TRITON_FRONTEND_PATH_ENV = "CPPMEGA_MLX_TRITON_FRONTEND_PATH"
 FLA_SOURCE_PATH_ENV = "CPPMEGA_MLX_FLA_SOURCE_PATH"
+TRITON_FRONTEND_UNSAFE_IMPORT_ENV = "CPPMEGA_V4_ENABLE_UNSAFE_TRITON_IMPORT"
 
 _TRITON_FRONTEND_ROOTS = (
     Path("/Users/dave/sources/tilelang"),
@@ -62,9 +63,47 @@ def ensure_fla_root() -> str | None:
     return None
 
 
+def unsafe_triton_frontend_import_enabled() -> bool:
+    """Return True only when the caller explicitly accepts unsafe Triton imports.
+
+    Some local Triton checkouts abort the Python process during import instead
+    of raising a Python exception. Path D status probes run in normal test and
+    benchmark discovery, so they must fail closed unless the developer opts in.
+    """
+
+    return os.environ.get(TRITON_FRONTEND_UNSAFE_IMPORT_ENV, "").strip().lower() in {
+        "1",
+        "true",
+        "yes",
+        "on",
+    }
+
+
+def unsafe_triton_frontend_import_disabled_reason(root: str | None) -> str:
+    root_text = f" root={root}" if root else " root=<not found>"
+    return (
+        "triton frontend not importable by default: unsafe import disabled "
+        f"({TRITON_FRONTEND_UNSAFE_IMPORT_ENV}=1 required); Path D runtime "
+        f"adapter not reached;{root_text}"
+    )
+
+
+def unsafe_fla_import_disabled_reason(root: str | None) -> str:
+    root_text = f" root={root}" if root else " root=<not found>"
+    return (
+        "FLA not importable by default: unsafe Path D import disabled "
+        f"({TRITON_FRONTEND_UNSAFE_IMPORT_ENV}=1 required); Path D runtime "
+        f"adapter not reached;{root_text}"
+    )
+
+
 __all__ = [
     "FLA_SOURCE_PATH_ENV",
     "TRITON_FRONTEND_PATH_ENV",
+    "TRITON_FRONTEND_UNSAFE_IMPORT_ENV",
     "ensure_fla_root",
     "ensure_triton_frontend_root",
+    "unsafe_fla_import_disabled_reason",
+    "unsafe_triton_frontend_import_disabled_reason",
+    "unsafe_triton_frontend_import_enabled",
 ]