Path C: sync sparse MLA static ABI inputs into banks

Author: Codex

cppmega_mlx/models/hybrid_lm.py

@@ -1272,6 +1272,58 @@ def path_c_fused_in_region_parameter_bank_aliases(
                 break
         return out
 
+    def _path_c_static_real_abi_input_values(
+        self,
+        abi_map: Mapping[str, Any],
+    ) -> dict[str, mx.array]:
+        """Return small non-trainable real-ABI inputs required by Path C.
+
+        The fused train-block ABI contains a few real inputs that are neither
+        trainable parameters nor per-batch runtime tensors. For Sparse-MLA
+        these are the softmax scale and optional attention sinks. They must be
+        written into model-owned banks along with parameter values; leaving the
+        zero-initialized bank slots in place makes the forward/backward kernel
+        run with ``sm_scale=0`` and kills q/kv gradients.
+        """
+
+        attention = self.config.attention_config("dsa")
+        sm_scale = float(attention.q_head_dim) ** -0.5
+        values: dict[str, mx.array] = {}
+        for logical_name, info in abi_map.items():
+            if not isinstance(info, Mapping):
+                continue
+            shape = tuple(
+                int(dim)
+                for dim in tuple(info.get("logical_shape", info.get("shape", ())))
+            )
+            if str(logical_name).endswith("sparse_mla_sm_scale"):
+                values[str(logical_name)] = mx.array(
+                    [sm_scale], dtype=mx.float32
+                ).reshape(shape or (1,))
+            elif str(logical_name).endswith("sparse_mla_sinks"):
+                values[str(logical_name)] = mx.zeros(shape, dtype=mx.float32)
+            elif str(logical_name).endswith("sparse_mla_has_sinks"):
+                values[str(logical_name)] = mx.zeros(shape, dtype=mx.int32)
+        return values
+
+    def _sync_path_c_static_real_abi_inputs_into_bank(
+        self,
+        *,
+        abi_map: Mapping[str, Any],
+        buffers: Mapping[str, Any],
+    ) -> tuple[list[str], list[tuple[str, str]]]:
+        synced: list[str] = []
+        skipped: list[tuple[str, str]] = []
+        for logical_name, value in sorted(
+            self._path_c_static_real_abi_input_values(abi_map).items()
+        ):
+            try:
+                write_into_bank_slot(abi_map, buffers, logical_name, value)
+                synced.append(logical_name)
+            except Exception as exc:
+                skipped.append((logical_name, f"{type(exc).__name__}: {exc}"))
+        return synced, skipped
+
     def _path_c_lookup_parameter_holder(
         self,
         parameter_name: str,
@@ -1378,6 +1430,13 @@ def sync_path_c_in_region_parameters_into_bank(
         )
         synced: list[str] = []
         skipped: list[tuple[str, str]] = []
+        static_synced, static_skipped = (
+            self._sync_path_c_static_real_abi_inputs_into_bank(
+                abi_map=abi_map,
+                buffers=buffers,
+            )
+        )
+        skipped.extend(static_skipped)
         for parameter_name, info in sorted(aliases.items()):
             tensor = self._path_c_get_parameter_tensor(parameter_name)
             if tensor is None:
@@ -1414,6 +1473,11 @@ def sync_path_c_in_region_parameters_into_bank(
                 {"parameter_name": name, "reason": reason}
                 for name, reason in skipped
             ],
+            "static_real_abi_inputs_synced": static_synced,
+            "static_real_abi_inputs_skipped": [
+                {"logical_name": name, "reason": reason}
+                for name, reason in static_skipped
+            ],
         }
 
     def bind_path_c_in_region_parameter_views_into_bank(
@@ -1486,6 +1550,13 @@ def bind_path_c_in_region_parameter_views_into_bank(
         )
         bound: list[str] = []
         skipped: list[tuple[str, str]] = []
+        static_synced, static_skipped = (
+            self._sync_path_c_static_real_abi_inputs_into_bank(
+                abi_map=abi_map,
+                buffers=buffers,
+            )
+        )
+        skipped.extend(static_skipped)
         for parameter_name, info in sorted(aliases.items()):
             tensor = self._path_c_get_parameter_tensor(parameter_name)
             if tensor is None:
@@ -1524,6 +1595,11 @@ def bind_path_c_in_region_parameter_views_into_bank(
             "in_region_parameter_count": len(bound),
             "in_region_parameter_names": tuple(sorted(bound)),
             "in_region_parameter_bank_aliases": dict(aliases),
+            "static_real_abi_inputs_synced": static_synced,
+            "static_real_abi_inputs_skipped": [
+                {"logical_name": name, "reason": reason}
+                for name, reason in static_skipped
+            ],
         }
 
     def path_c_fused_first_in_region_layer_index(

tests/test_hybrid_lm_path_c_physical_abi_bank_owner.py

@@ -9,6 +9,8 @@
 
 from __future__ import annotations
 
+from typing import Any
+
 import mlx.core as mx
 import pytest
 
@@ -188,7 +190,7 @@ def test_bind_path_c_in_region_parameter_views_into_bank_replaces_attributes(
     bank = bank_owner.buffers[sample_info["bank"]]
     slot = bank[sample_info["offset"] : sample_info["offset"] + sample_info["size"]]
     parts = sample_small.split(".")
-    holder: object = model
+    holder: Any = model
     for part in parts[:-1]:
         holder = holder[int(part)] if part.isdigit() else getattr(holder, part)
     param_value = getattr(holder, parts[-1])
@@ -226,3 +228,54 @@ def test_sync_path_c_in_region_parameters_into_bank_updates_bank_slots(
     bank = bank_owner.buffers[bank_name]
     slot = bank[info["offset"] : info["offset"] + info["size"]]
     assert mx.allclose(slot, sentinel).item() is True
+
+
+def test_bind_path_c_syncs_sparse_mla_static_real_abi_inputs() -> None:
+    model = build_local_gb10_quarter_tiny_smoke_model()
+    sequence_length = 127
+    bank_owner = model.make_path_c_physical_abi_bank_owner(
+        sequence_length=sequence_length,
+    )
+    assert bank_owner is not None
+    report = model.bind_path_c_in_region_parameter_views_into_bank(
+        bank_owner,
+        sequence_length=sequence_length,
+    )
+    prim_func = model.path_c_fused_train_block_prim_func(
+        sequence_length=sequence_length,
+    )
+    assert prim_func is not None
+    abi_map = dict(
+        getattr(prim_func, "_cppmega_path_c_physical_buffer_abi_map", {})
+        or {}
+    )
+
+    def bank_slot(logical_name: str) -> mx.array:
+        info = abi_map[logical_name]
+        bank = bank_owner.buffers[str(info["bank"])]
+        offset = int(info["offset"])
+        size = int(info["size"])
+        return bank[offset : offset + size]
+
+    sm_scale_name = next(
+        name for name in abi_map if name.endswith("sparse_mla_sm_scale")
+    )
+    sinks_name = next(
+        name for name in abi_map if name.endswith("sparse_mla_sinks")
+    )
+    has_sinks_name = next(
+        name for name in abi_map if name.endswith("sparse_mla_has_sinks")
+    )
+    sm_scale = bank_slot(sm_scale_name)
+    sinks = bank_slot(sinks_name)
+    has_sinks = bank_slot(has_sinks_name)
+    mx.eval(sm_scale, sinks, has_sinks)
+
+    assert report["static_real_abi_inputs_synced"] == [
+        has_sinks_name,
+        sinks_name,
+        sm_scale_name,
+    ]
+    assert sm_scale.tolist() == [0.5]
+    assert sinks.tolist() == [0.0, 0.0, 0.0, 0.0]
+    assert has_sinks.tolist() == [0]