HybridTinyLM: own Path C physical ABI banks so fused train-block route binds

The m04 fp8_path_c training route's auto-install path is gated on the model
exposing a callable make_path_c_physical_abi_bank_owner(sequence_length=...)
factory. Without one, the route reports
FP8_PATH_C_FUSED_TRAIN_BLOCK_BANKS_MISSING_STATUS, the artifact never gets
wrapped into a PathCFusedTrainBlockTrainingRuntime, and every cell in the
1B Path C matrix records 'fused_train_block_runtime_available=False'.

This change adds two model methods:
  - path_c_fused_train_block_prim_func(...) -> PrimFunc | None
    Drives the existing route discovery + plan_path_c_fusion_schedule_for_region
    pipeline and returns the generated PrimFunc that carries the physical-ABI
    attrs. Cached by caller only; no implicit reuse.
  - make_path_c_physical_abi_bank_owner(...) -> PathCPhysicalAbiBankOwner | None
    Reads _cppmega_path_c_physical_buffer_abi_{map,shapes} off the prim_func,
    allocates each declared bank as a zero-init mx.array sized to the exact
    physical-ABI dtype/shape, and returns a validated owner via
    make_physical_abi_bank_owner (no hidden packing, no implicit copy).

After this commit, fp8_path_c_training_route_payload_for_model on a tiny
HybridTinyLM reports binding.status='ok', physical_abi_binding_ready=True,
fused_artifact_bound=True, and runtime_uses_fused_train_block=True with all
three banks (float32, uint8, int32) bound to the model owner.

Behaviour change checked into the existing
test_fp8_path_c_training_route_reports_model_owned_physical_bank_plan: the
binding status assertion flipped from BANKS_MISSING to 'ok'; bank-plan
content invariants unchanged.

Verified: pytest -q on every existing path_c m04 test plus the new
tests/test_hybrid_lm_path_c_physical_abi_bank_owner.py (4 cases) all pass;
ruff and pyright clean on the touched module.

Author: apstenku123

cppmega_mlx/models/hybrid_lm.py

@@ -43,7 +43,11 @@
 from cppmega_mlx.nn.structure_embedding import CppMegaStructureEmbedding
 from cppmega_mlx.recipes.pattern import ExpandedNamPattern, NamLayer, expand_nam_pattern
 from cppmega_mlx.runtime.kernel_policy import KernelPath, selected_path
-from cppmega_mlx.runtime.path_c_physical_abi import PathCLogicalBufferOwner
+from cppmega_mlx.runtime.path_c_physical_abi import (
+    PathCLogicalBufferOwner,
+    PathCPhysicalAbiBankOwner,
+    make_physical_abi_bank_owner,
+)
 from cppmega_mlx.runtime.path_c_taps import emit_and_tap_path_c_tensor
 from cppmega_mlx.training.mtp import MinimalMTPHead, MTPLossConfig
 
@@ -74,6 +78,33 @@
 
 HybridAttentionMode = Literal["mla", "dsa", "full", "gqa"]
 
+_PATH_C_BANK_DTYPES: dict[str, mx.Dtype] = {
+    "bool": mx.bool_,
+    "uint8": mx.uint8,
+    "int8": mx.int8,
+    "uint16": mx.uint16,
+    "int16": mx.int16,
+    "uint32": mx.uint32,
+    "int32": mx.int32,
+    "uint64": mx.uint64,
+    "int64": mx.int64,
+    "float16": mx.float16,
+    "bfloat16": mx.bfloat16,
+    "float32": mx.float32,
+}
+
+
+def _path_c_bank_dtype(name: str) -> mx.Dtype:
+    """Translate a Path C physical-ABI dtype string into an ``mx.Dtype``."""
+    try:
+        return _PATH_C_BANK_DTYPES[name]
+    except KeyError as exc:
+        raise ValueError(
+            f"unsupported Path C physical ABI bank dtype {name!r}"
+        ) from exc
+
+
+
 
 class StructureEmbeddingConfigKwargs(TypedDict):
     hidden_size: int
@@ -1061,6 +1092,110 @@ def path_c_fusion_regions(
             sequence_length=sequence_length,
         )
 
+    def path_c_fused_train_block_prim_func(
+        self,
+        *,
+        sequence_length: int | None = None,
+    ) -> Any | None:
+        """Materialise the fused-train-block PrimFunc for this model.
+
+        This drives the same schedule planner the production runtime uses and
+        returns the generated PrimFunc (carrying the physical-ABI attrs), or
+        ``None`` when no Path C route region exists. The model never caches the
+        artifact — callers that need to reuse it across calls should hold the
+        returned reference themselves.
+        """
+        from cppmega_mlx.runtime.path_c_fusion_schedules import (
+            plan_path_c_fusion_schedule_for_region,
+        )
+
+        regions = self.path_c_fusion_regions(
+            include_backward=True,
+            sequence_length=sequence_length,
+        )
+        if not regions:
+            return None
+        selected = max(
+            regions,
+            key=lambda region: (
+                len(region.nodes),
+                len(region.edges),
+                region.name,
+            ),
+        )
+        planned = plan_path_c_fusion_schedule_for_region(
+            selected,
+            include_backward=True,
+        )
+        schedule_target = getattr(planned, "schedule_target", None)
+        if schedule_target is None:
+            return None
+        return schedule_target.schedule_template(planned.region)
+
+    def make_path_c_physical_abi_bank_owner(
+        self,
+        *,
+        sequence_length: int | None = None,
+    ) -> PathCPhysicalAbiBankOwner | None:
+        """Allocate the physical ABI banks the generated fused PrimFunc needs.
+
+        Returns a validated :class:`PathCPhysicalAbiBankOwner` whose ``buffers``
+        are freshly zero-initialised MLX arrays sized exactly to the generated
+        ``_cppmega_path_c_physical_buffer_abi_shapes`` map (dtype is taken from
+        the corresponding logical-buffer placement so every bank reflects the
+        generated kernel ABI). The owner never repacks or copies model tensors;
+        it only owns the bank arrays so the runtime can bind kernel arguments
+        in declared order.
+
+        Returns ``None`` when no Path C route region exists for the model.
+        """
+        prim_func = self.path_c_fused_train_block_prim_func(
+            sequence_length=sequence_length,
+        )
+        if prim_func is None:
+            return None
+        abi_map = dict(
+            getattr(prim_func, "_cppmega_path_c_physical_buffer_abi_map", {})
+            or {}
+        )
+        abi_shapes = dict(
+            getattr(prim_func, "_cppmega_path_c_physical_buffer_abi_shapes", {})
+            or {}
+        )
+        if not abi_map or not abi_shapes:
+            return None
+        bank_dtypes: dict[str, str] = {}
+        for placement in abi_map.values():
+            bank = str(placement["bank"])
+            dtype = str(placement["dtype"])
+            existing = bank_dtypes.setdefault(bank, dtype)
+            if existing != dtype:
+                raise ValueError(
+                    f"conflicting bank dtype for {bank!r}: "
+                    f"{existing!r} vs {dtype!r}"
+                )
+        bank_buffers: dict[str, mx.array] = {}
+        for bank, shape in abi_shapes.items():
+            dtype_name = bank_dtypes.get(str(bank))
+            if dtype_name is None:
+                raise ValueError(
+                    f"no logical buffer is placed inside physical bank {bank!r}"
+                )
+            mx_dtype = _path_c_bank_dtype(dtype_name)
+            bank_buffers[str(bank)] = mx.zeros(
+                tuple(int(dim) for dim in tuple(shape)),
+                dtype=mx_dtype,
+            )
+        profile_name = str(
+            getattr(self, "path_c_profile_name", "HybridTinyLM")
+        )
+        return make_physical_abi_bank_owner(
+            f"{profile_name}.path_c_physical_abi_banks",
+            abi_map,
+            abi_shapes,
+            bank_buffers,
+        )
+
     def __call__(
         self,
         input_ids: mx.array,

tests/test_hybrid_lm_path_c_physical_abi_bank_owner.py

@@ -0,0 +1,100 @@
+"""HybridTinyLM.make_path_c_physical_abi_bank_owner.
+
+Verifies that the model exposes the model-owned physical ABI bank factory
+the m04 Path C auto-install path requires. The check exercises the actual
+production schedule (no monkeypatching) and asserts that the resulting
+owner satisfies the structural contract enforced by
+``cppmega_mlx.runtime.path_c_physical_abi.make_physical_abi_bank_owner``.
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+import pytest
+
+from cppmega_mlx.recipes.model_factory import (
+    build_local_gb10_quarter_tiny_smoke_model,
+)
+from cppmega_mlx.runtime.path_c_physical_abi import (
+    PathCPhysicalAbiBankOwner,
+)
+
+
+@pytest.fixture(scope="module")
+def model():
+    return build_local_gb10_quarter_tiny_smoke_model()
+
+
+def test_make_path_c_physical_abi_bank_owner_returns_validated_owner(model) -> None:
+    owner = model.make_path_c_physical_abi_bank_owner(sequence_length=64)
+    assert isinstance(owner, PathCPhysicalAbiBankOwner)
+    assert owner.owner_name == "local_gb10_quarter.path_c_physical_abi_banks"
+    assert owner.binding_payload["status"] == "ok"
+    assert owner.hidden_packing_performed is False
+    assert owner.no_hidden_allocation_policy is True
+    # Bank specs must be non-empty and every bank buffer must already be
+    # allocated as an mx.array with the spec's shape and dtype.
+    assert owner.bank_specs, "expected at least one physical ABI bank"
+    required = set(owner.required_bank_buffers)
+    assert set(owner.buffers) == required
+    for spec in owner.bank_specs:
+        buf = owner.buffers[spec.name]
+        assert isinstance(buf, mx.array)
+        assert tuple(buf.shape) == spec.shape
+        assert str(buf.dtype).rsplit(".", 1)[-1] == spec.dtype
+
+
+def test_make_path_c_physical_abi_bank_owner_buffers_are_zero_initialised(model) -> None:
+    owner = model.make_path_c_physical_abi_bank_owner(sequence_length=64)
+    for spec in owner.bank_specs:
+        buf = owner.buffers[spec.name]
+        sample = buf.astype(mx.float32)
+        mx.eval(sample)
+        assert float(sample.sum().item()) == 0.0
+
+
+def test_path_c_fused_train_block_prim_func_exposes_physical_abi_attrs(model) -> None:
+    prim_func = model.path_c_fused_train_block_prim_func(sequence_length=64)
+    assert prim_func is not None
+    abi_map = dict(
+        getattr(prim_func, "_cppmega_path_c_physical_buffer_abi_map", {}) or {}
+    )
+    abi_shapes = dict(
+        getattr(prim_func, "_cppmega_path_c_physical_buffer_abi_shapes", {}) or {}
+    )
+    # The route region for local_gb10_quarter must emit a banked physical ABI.
+    assert abi_map, "physical ABI map must not be empty"
+    assert abi_shapes, "physical ABI bank shapes must not be empty"
+    # Bank names referenced by the placements must all appear in the shape map.
+    placement_banks = {str(info["bank"]) for info in abi_map.values()}
+    assert placement_banks <= set(abi_shapes)
+
+
+def test_owner_unblocks_m04_fp8_path_c_runtime_binding(model) -> None:
+    """The owner must let the m04 route_payload reach a bound runtime."""
+
+    import importlib.util
+    from pathlib import Path
+
+    spec = importlib.util.spec_from_file_location(
+        "m04_train_step", str(Path("scripts/m04_train_step.py").resolve())
+    )
+    assert spec is not None and spec.loader is not None
+    m04 = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(m04)
+
+    args = m04.build_parser().parse_args([
+        "--synthetic", "--dtype", "fp8_path_c",
+        "--model-profile", "local_gb10_quarter",
+        "--dry-run-json", "--output", "/tmp/test_owner_route.json",
+        "--data-path", "/dev/null", "--data-format", "npz",
+    ])
+    route = m04.fp8_path_c_training_route_payload_for_model(args, model)
+    binding = route.get("path_c_fusion", {}).get("runtime_training_binding", {})
+    assert binding.get("runtime_uses_fused_train_block") is True, route
+    assert binding.get("physical_abi_binding_ready") is True, binding
+    assert binding.get("fused_artifact_bound") is True, binding
+    assert binding.get("missing_bank_buffers") == [], binding
+    assert binding.get("bank_buffer_owner") == (
+        "local_gb10_quarter.path_c_physical_abi_banks"
+    )

tests/test_m04_train_step.py

@@ -3871,9 +3871,15 @@ def test_fp8_path_c_training_route_reports_model_owned_physical_bank_plan(
     )
 
     binding = route["path_c_fusion"]["runtime_training_binding"]
-    assert binding["status"] == (
-        m04_train_step.FP8_PATH_C_FUSED_TRAIN_BLOCK_BANKS_MISSING_STATUS
-    )
+    # HybridTinyLM now provides ``make_path_c_physical_abi_bank_owner`` so
+    # the production auto-install path resolves the banks before the route
+    # payload is computed; the runtime training binding reports ``ok``.
+    assert binding["status"] == "ok", binding
+    assert binding["runtime_uses_fused_train_block"] is True
+    assert binding["physical_abi_binding_ready"] is True
+    assert binding["fused_artifact_bound"] is True
+    assert binding["missing_bank_buffers"] == []
+    assert binding["provided_bank_buffers"] == binding["required_bank_buffers"]
     bank_plan = binding["model_owned_physical_abi_bank_plan"]
     assert bank_plan["status"] == "model_owned_physical_abi_banks_required"
     assert bank_plan["owner_attribute"] == "path_c_physical_abi_bank_owner"